add partial sort for gslice

Author: hay.123

gslice/gslice.go

@@ -1011,6 +1011,11 @@ func StableSortBy[T any](s []T, less func(T, T) bool) {
 	_ = iter.ToSlice(iter.StableSortBy(less, iter.StealSlice(s)))
 }
 
+// PartialSort
+func PartialSort[T constraints.Ordered](s []T, k int) {
+	_ = iter.ToSlice(iter.PartialSort(k, iter.StealSlice(s)))
+}
+
 // TypeAssert converts a slice from type From to type To by type assertion.
 //
 // 🚀 EXAMPLE:

gslice/gslice_test.go

@@ -623,6 +623,50 @@ 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
+	}
+
+	// Already sorted
+	{
+		s := []int{1, 2, 3, 4, 5}
+		PartialSort(s, 3)
+		assert.Equal(t, []int{1, 2, 3, 4, 5}, s) // Unchanged
+	}
+}
+
 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

@@ -48,3 +48,32 @@ func Sort[T constraints.Ordered](v []T) {
 		siftDown(v[:i], 0)
 	}
 }
+
+func PartialSort[T constraints.Ordered](v []T, k int) {
+	n := len(v)
+
+	if k <= 0 {
+		return
+	}
+
+	if k >= n {
+		Sort(v)
+		return
+	}
+
+	// Build a max-heap from the first k elements
+	for j := (k - 1) / 2; j >= 0; j-- {
+		siftDown(v[:k], j)
+	}
+
+	// Iterate through the rest of the slice
+	for j := k; j < n; j++ {
+		if v[j] < v[0] {
+			v[0], v[j] = v[j], v[0]
+			siftDown(v[:k], 0)
+		}
+	}
+
+	// Sort the heap to get the final k smallest elements in order
+	Sort(v[:k])
+}

internal/iter/operations.go

@@ -923,6 +923,16 @@ 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)
+}
+
 // Contains returns whether the element occur in iterator.
 func Contains[T comparable](v T, i Iter[T]) bool {
 	for _, vv := range i.Next(ALL) {