From 32cb0782991d6b35ea7b33712b1892a6134484d6 Mon Sep 17 00:00:00 2001 From: Mark Grayson <118820911+sahildando@users.noreply.github.com> Date: Mon, 13 Apr 2026 01:12:43 +0530 Subject: [PATCH 1/2] Optimize SelectK with introspective quickselect --- search/doc.go | 6 ++- search/selectk.go | 100 +++++++++++++++++++++++++++++++++-------- search/selectk_test.go | 5 ++- 3 files changed, 91 insertions(+), 20 deletions(-) diff --git a/search/doc.go b/search/doc.go index 4c71e36b4..cd56ca0d1 100644 --- a/search/doc.go +++ b/search/doc.go @@ -1,2 +1,6 @@ -// Package search is a subpackage dedicated to all searching algorithms related to slices/arrays. +// Package search contains classic and practical search algorithms for arrays and slices. +// +// The implementations in this package are designed to be easy to read first, +// while still reflecting good algorithmic choices (time/space complexity, +// boundary validation, and benchmark coverage). package search diff --git a/search/selectk.go b/search/selectk.go index bb81937b4..e64bfb6a9 100644 --- a/search/selectk.go +++ b/search/selectk.go @@ -1,32 +1,96 @@ package search +import ( + "math/bits" + "sort" +) + +// SelectK returns the k-th largest element in array. +// +// Time complexity is expected O(n) thanks to quickselect partitioning. +// A depth limit is applied and falls back to sorting the narrowed range, +// which guarantees O(n log n) worst-case behavior. +// +// The function mutates the input slice in-place. func SelectK(array []int, k int) (int, error) { - if k > len(array) { + n := len(array) + if n == 0 || k < 1 || k > n { return -1, ErrNotFound } - return selectK(array, 0, len(array), len(array)-k), nil + + // k-th largest -> index in zero-based ascending order. + idx := n - k + return selectK(array, idx), nil } -// search the element which index is idx -func selectK(array []int, l, r, idx int) int { - index := partition(array, l, r) - if index == idx { - return array[index] +// selectK returns the element that would appear at idx in sorted order. +func selectK(array []int, idx int) int { + l, r := 0, len(array) + depthLimit := 2 * bits.Len(uint(len(array))) + + for r-l > 1 { + if depthLimit == 0 { + sort.Ints(array[l:r]) + return array[idx] + } + depthLimit-- + + leftPivot, rightPivot := partition3(array, l, r) + switch { + case idx < leftPivot: + r = leftPivot + case idx >= rightPivot: + l = rightPivot + default: + return array[idx] + } } - if index < idx { - return selectK(array, index+1, r, idx) + + return array[l] +} + +// partition3 applies a Dutch National Flag partition around a pivot and +// returns [leftPivot, rightPivot), the range that equals the pivot. +func partition3(array []int, l, r int) (leftPivot, rightPivot int) { + pivotIdx := medianOfThreeIndex(array, l, l+(r-l)/2, r-1) + pivot := array[pivotIdx] + array[l], array[pivotIdx] = array[pivotIdx], array[l] + + lt, i, gt := l, l+1, r + for i < gt { + switch { + case array[i] < pivot: + lt++ + array[i], array[lt] = array[lt], array[i] + i++ + case array[i] > pivot: + gt-- + array[i], array[gt] = array[gt], array[i] + default: + i++ + } } - return selectK(array, l, index, idx) + + array[l], array[lt] = array[lt], array[l] + return lt, gt } -func partition(array []int, l, r int) int { - elem, j := array[l], l+1 - for i := l + 1; i < r; i++ { - if array[i] <= elem { - array[i], array[j] = array[j], array[i] - j++ +func medianOfThreeIndex(array []int, a, b, c int) int { + if array[a] < array[b] { + if array[b] < array[c] { + return b } + if array[a] < array[c] { + return c + } + return a + } + + if array[a] < array[c] { + return a + } + if array[b] < array[c] { + return c } - array[l], array[j-1] = array[j-1], array[l] - return j - 1 + return b } diff --git a/search/selectk_test.go b/search/selectk_test.go index bf81fb9d1..14686b476 100644 --- a/search/selectk_test.go +++ b/search/selectk_test.go @@ -15,7 +15,10 @@ func TestSelectK(t *testing.T) { {[]int{1, 2, 3, 4, 5}, 1, 5, nil, "sorted data"}, {[]int{5, 4, 3, 2, 1}, 2, 4, nil, "reversed data"}, {[]int{3, 1, 2, 5, 4}, 3, 3, nil, "random data"}, - {[]int{3, 2, 1, 5, 4}, 10, -1, ErrNotFound, " absent data"}, + {[]int{3, 2, 1, 5, 4}, 10, -1, ErrNotFound, "absent data"}, + {[]int{3, 2, 1, 5, 4}, 0, -1, ErrNotFound, "invalid zero k"}, + {[]int{}, 1, -1, ErrNotFound, "empty data"}, + {[]int{5, 5, 5, 5, 5}, 3, 5, nil, "duplicate data"}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { From 113baccf17c19a34153b5f89d61472afe1c5cbd4 Mon Sep 17 00:00:00 2001 From: Mark Grayson <118820911+sahildando@users.noreply.github.com> Date: Mon, 13 Apr 2026 01:14:17 +0530 Subject: [PATCH 2/2] Revert "Optimize SelectK with introspective quickselect and robust partitioning" --- search/doc.go | 6 +-- search/selectk.go | 100 ++++++++--------------------------------- search/selectk_test.go | 5 +-- 3 files changed, 20 insertions(+), 91 deletions(-) diff --git a/search/doc.go b/search/doc.go index cd56ca0d1..4c71e36b4 100644 --- a/search/doc.go +++ b/search/doc.go @@ -1,6 +1,2 @@ -// Package search contains classic and practical search algorithms for arrays and slices. -// -// The implementations in this package are designed to be easy to read first, -// while still reflecting good algorithmic choices (time/space complexity, -// boundary validation, and benchmark coverage). +// Package search is a subpackage dedicated to all searching algorithms related to slices/arrays. package search diff --git a/search/selectk.go b/search/selectk.go index e64bfb6a9..bb81937b4 100644 --- a/search/selectk.go +++ b/search/selectk.go @@ -1,96 +1,32 @@ package search -import ( - "math/bits" - "sort" -) - -// SelectK returns the k-th largest element in array. -// -// Time complexity is expected O(n) thanks to quickselect partitioning. -// A depth limit is applied and falls back to sorting the narrowed range, -// which guarantees O(n log n) worst-case behavior. -// -// The function mutates the input slice in-place. func SelectK(array []int, k int) (int, error) { - n := len(array) - if n == 0 || k < 1 || k > n { + if k > len(array) { return -1, ErrNotFound } - - // k-th largest -> index in zero-based ascending order. - idx := n - k - return selectK(array, idx), nil + return selectK(array, 0, len(array), len(array)-k), nil } -// selectK returns the element that would appear at idx in sorted order. -func selectK(array []int, idx int) int { - l, r := 0, len(array) - depthLimit := 2 * bits.Len(uint(len(array))) - - for r-l > 1 { - if depthLimit == 0 { - sort.Ints(array[l:r]) - return array[idx] - } - depthLimit-- - - leftPivot, rightPivot := partition3(array, l, r) - switch { - case idx < leftPivot: - r = leftPivot - case idx >= rightPivot: - l = rightPivot - default: - return array[idx] - } +// search the element which index is idx +func selectK(array []int, l, r, idx int) int { + index := partition(array, l, r) + if index == idx { + return array[index] } - - return array[l] -} - -// partition3 applies a Dutch National Flag partition around a pivot and -// returns [leftPivot, rightPivot), the range that equals the pivot. -func partition3(array []int, l, r int) (leftPivot, rightPivot int) { - pivotIdx := medianOfThreeIndex(array, l, l+(r-l)/2, r-1) - pivot := array[pivotIdx] - array[l], array[pivotIdx] = array[pivotIdx], array[l] - - lt, i, gt := l, l+1, r - for i < gt { - switch { - case array[i] < pivot: - lt++ - array[i], array[lt] = array[lt], array[i] - i++ - case array[i] > pivot: - gt-- - array[i], array[gt] = array[gt], array[i] - default: - i++ - } + if index < idx { + return selectK(array, index+1, r, idx) } - - array[l], array[lt] = array[lt], array[l] - return lt, gt + return selectK(array, l, index, idx) } -func medianOfThreeIndex(array []int, a, b, c int) int { - if array[a] < array[b] { - if array[b] < array[c] { - return b +func partition(array []int, l, r int) int { + elem, j := array[l], l+1 + for i := l + 1; i < r; i++ { + if array[i] <= elem { + array[i], array[j] = array[j], array[i] + j++ } - if array[a] < array[c] { - return c - } - return a - } - - if array[a] < array[c] { - return a - } - if array[b] < array[c] { - return c } - return b + array[l], array[j-1] = array[j-1], array[l] + return j - 1 } diff --git a/search/selectk_test.go b/search/selectk_test.go index 14686b476..bf81fb9d1 100644 --- a/search/selectk_test.go +++ b/search/selectk_test.go @@ -15,10 +15,7 @@ func TestSelectK(t *testing.T) { {[]int{1, 2, 3, 4, 5}, 1, 5, nil, "sorted data"}, {[]int{5, 4, 3, 2, 1}, 2, 4, nil, "reversed data"}, {[]int{3, 1, 2, 5, 4}, 3, 3, nil, "random data"}, - {[]int{3, 2, 1, 5, 4}, 10, -1, ErrNotFound, "absent data"}, - {[]int{3, 2, 1, 5, 4}, 0, -1, ErrNotFound, "invalid zero k"}, - {[]int{}, 1, -1, ErrNotFound, "empty data"}, - {[]int{5, 5, 5, 5, 5}, 3, 5, nil, "duplicate data"}, + {[]int{3, 2, 1, 5, 4}, 10, -1, ErrNotFound, " absent data"}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) {