diff --git a/gslice/gslice.go b/gslice/gslice.go index d213a89..84e6627 100644 --- a/gslice/gslice.go +++ b/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: diff --git a/gslice/gslice_test.go b/gslice/gslice_test.go index d5c9a6d..04719c1 100644 --- a/gslice/gslice_test.go +++ b/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})) diff --git a/internal/heapsort/sort.go b/internal/heapsort/sort.go index a9b2caa..6cc756a 100644 --- a/internal/heapsort/sort.go +++ b/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]) }) } diff --git a/internal/iter/operations.go b/internal/iter/operations.go index 0ee8897..3a3da2c 100644 --- a/internal/iter/operations.go +++ b/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) {