feat(gslice): Add partial sort for gslice

gslice/gslice.go

@@ -1011,6 +1011,17 @@ func StableSortBy[T any](s []T, less func(T, T) bool) {
 	_ = iter.ToSlice(iter.StableSortBy(less, iter.StealSlice(s)))
 }
 
+// PartialSort sorts the first k smallest elements in order, with the remaining elements left unordered
+func PartialSort[T constraints.Ordered](s []T, k int) {
+	_ = iter.ToSlice(iter.PartialSort(k, iter.StealSlice(s)))
+}
+
+// PartialSortBy reorders the slice so that the first k elements are the top-k
+// according to the comparison function `less`, and sorted in that order.
+func PartialSortBy[T constraints.Ordered](s []T, k int, less func(T, T) bool) {
+	_ = iter.ToSlice(iter.PartialSortBy(k, less, iter.StealSlice(s)))
+}
+
 // TypeAssert converts a slice from type From to type To by type assertion.
 //
 // 🚀 EXAMPLE:

gslice/gslice_test.go

@@ -623,6 +623,81 @@ func TestStableSortBy(t *testing.T) {
 	}
 }
 
+func TestPartialSort(t *testing.T) {
+	// Basic case - first k elements sorted and smallest
+	{
+		s := []int{5, 2, 9, 1, 5, 6}
+		PartialSort(s, 3)
+		assert.Equal(t, []int{1, 2, 5, 9, 5, 6}, s) // First 3 are smallest + sorted
+	}
+
+	// k == length (full sort)
+	{
+		s := []int{3, 1, 4}
+		PartialSort(s, 3)
+		assert.Equal(t, []int{1, 3, 4}, s) // Fully sorted
+	}
+
+	// k > length (behaves like full sort)
+	{
+		s := []int{7, 3}
+		PartialSort(s, 5)
+		assert.Equal(t, []int{3, 7}, s) // Treats as full sort
+	}
+
+	// Empty slice
+	{
+		s := []int{}
+		PartialSort(s, 2)
+		assert.Equal(t, []int{}, s) // No panic
+	}
+
+	// k == 0 (no-op)
+	{
+		s := []int{5, 2, 8}
+		PartialSort(s, 0)
+		assert.Equal(t, []int{5, 2, 8}, s) // Unmodified
+	}
+
+}
+
+func TestPartialSortByDescending(t *testing.T) {
+	// Basic case - first k elements should be k largest, sorted descending
+	{
+		s := []int{5, 2, 9, 1, 5, 6}
+		PartialSortBy(s, 3, gvalue.Greater[int])
+		assert.Equal(t, []int{9, 6, 5, 1, 2, 5}, s) // First 3 are largest, sorted descending
+	}
+
+	// k == length (full descending sort)
+	{
+		s := []int{3, 1, 4}
+		PartialSortBy(s, 3, gvalue.Greater[int])
+		assert.Equal(t, []int{4, 3, 1}, s)
+	}
+
+	// k > length (full descending sort)
+	{
+		s := []int{3, 7}
+		PartialSortBy(s, 5, gvalue.Greater[int])
+		assert.Equal(t, []int{7, 3}, s)
+	}
+
+	// Empty slice
+	{
+		s := []int{}
+		PartialSortBy(s, 2, gvalue.Greater[int])
+		assert.Equal(t, []int{}, s)
+	}
+
+	// k == 0 (no-op)
+	{
+		s := []int{5, 2, 8}
+		PartialSortBy(s, 0, gvalue.Greater[int])
+		assert.Equal(t, []int{5, 2, 8}, s)
+	}
+}
+
 func TestTypeAssert(t *testing.T) {
 	assert.Equal(t, []int{1, 2, 3, 4}, TypeAssert[int, any]([]any{1, 2, 3, 4}))
 	assert.Equal(t, []any{1, 2, 3, 4}, TypeAssert[any, int]([]int{1, 2, 3, 4}))

internal/heapsort/sort.go

@@ -18,33 +18,83 @@ import (
 	"github.com/bytedance/gg/internal/constraints"
 )
 
-// siftDown implements the heap property on v[lo:hi].
-func siftDown[T constraints.Ordered](v []T, node int) {
+func siftDown[T constraints.Ordered](v []T, lo, hi, first int, less func(i, j int) bool) {
+	root := lo
 	for {
-		child := 2*node + 1
-		if child >= len(v) {
+		child := 2*root + 1
+		if child >= hi {
 			break
 		}
-		if child+1 < len(v) && v[child] < v[child+1] {
+		if child+1 < hi && less(first+child, first+child+1) {
 			child++
 		}
-		if v[node] >= v[child] {
+		if !less(first+root, first+child) {
 			return
 		}
-		v[node], v[child] = v[child], v[node]
-		node = child
+		v[first+root], v[first+child] = v[first+child], v[first+root]
+		root = child
 	}
 }
 
+func heapify[T constraints.Ordered](v []T, a, b int, less func(i, j int) bool) {
+	first := a
+	hi := b - a
+	for i := (hi - 1) / 2; i >= 0; i-- {
+		siftDown(v, i, hi, first, less)
+	}
+}
+
+func heapSort[T constraints.Ordered](v []T, a, b int, less func(i, j int) bool) {
+	first := a
+	lo := 0
+	hi := b - a
+
+	heapify(v, a, b, less)
+
+	for i := hi - 1; i >= 0; i-- {
+		v[first], v[first+i] = v[first+i], v[first]
+		siftDown(v, lo, i, first, less)
+	}
+}
+
+func partialSort[T constraints.Ordered](v []T, k int, less func(i, j int) bool) {
+	n := len(v)
+	if k <= 0 || n <= 1 {
+		return
+	}
+	if k >= n {
+		heapSort(v, 0, n, less)
+		return
+	}
+	heapify(v, 0, k, less)
+	for i := k; i < n; i++ {
+		if less(i, 0) {
+			v[0], v[i] = v[i], v[0]
+			siftDown(v, 0, k, 0, less)
+		}
+	}
+	heapSort(v, 0, k, less)
+}
+
 func Sort[T constraints.Ordered](v []T) {
-	// Build heap with the greatest element at the top.
-	for i := (len(v) - 1) / 2; i >= 0; i-- {
-		siftDown(v, i)
+	if len(v) <= 1 {
+		return
 	}
+	heapSort(v, 0, len(v), func(i, j int) bool { return v[i] < v[j] })
+}
+
+func PartialSort[T constraints.Ordered](v []T, k int) {
+	PartialSortBy(v, k, func(a, b T) bool { return a < b })
+}
 
-	// Pop elements into end of v.
-	for i := len(v) - 1; i >= 1; i-- {
-		v[0], v[i] = v[i], v[0]
-		siftDown(v[:i], 0)
+func PartialSortBy[T constraints.Ordered](v []T, k int, less func(a, b T) bool) {
+	n := len(v)
+	if k <= 0 || n <= 1 {
+		return
+	}
+	if k >= n {
+		heapSort(v, 0, n, func(i, j int) bool { return less(v[i], v[j]) })
+		return
 	}
+	partialSort(v, k, func(i, j int) bool { return less(v[i], v[j]) })
 }

internal/iter/operations.go

@@ -923,6 +923,25 @@ func Sort[T constraints.Ordered](i Iter[T]) Iter[T] {
 	return StealSlice(s)
 }
 
+// PartialSort sorts the first k smallest elements in ascending order, with the remaining elements left unordered,
+// and returns a new iterator.
+// - If k <= 0, returns the entire slice unmodified.
+// - If k >= len(slice), performs a full sort.
+func PartialSort[T constraints.Ordered](k int, iter Iter[T]) Iter[T] {
+	s := ToSlice(iter)
+	heapsort.PartialSort(s, k)
+	return StealSlice(s)
+}
+
+// PartialSortBy reorders the slice so that the first k elements are the top-k
+// according to the comparison function `less`, and sorted in that order.
+// return a new iterator.
+func PartialSortBy[T constraints.Ordered](k int, less func(T, T) bool, iter Iter[T]) Iter[T] {
+	s := ToSlice(iter)
+	heapsort.PartialSortBy(s, k, less)
+	return StealSlice(s)
+}
+
 // Contains returns whether the element occur in iterator.
 func Contains[T comparable](v T, i Iter[T]) bool {
 	for _, vv := range i.Next(ALL) {