Merge pull request #1 from sahildando/codex/optimize-code-for-performance-and-documentation

Optimize SelectK with introspective quickselect and robust partitioning

Author: sahildando

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

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
 }

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