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) {