82 lines
2.3 KiB
Go
82 lines
2.3 KiB
Go
package utils
|
|
|
|
import (
|
|
"cmp"
|
|
)
|
|
|
|
// NthElementSlice QuickSelect implementation
|
|
// k is the desired index value, where array[k] is the k+1 smallest element
|
|
// values between s[0, k-1] are guaranteed <= to s[k]
|
|
// values between s[k+1, len(s)] are guaranteed >= to s[k]
|
|
func NthElementSlice[S ~[]E, E cmp.Ordered](s S, k int) {
|
|
left := 0
|
|
right := len(s) - 1
|
|
for {
|
|
if left == right {
|
|
return
|
|
}
|
|
pivotIndex := (left + right) >> 1 // Use middle point as pivot. This could probably use random index
|
|
pivotIndex = partition(s, pivotIndex, left, right)
|
|
if k == pivotIndex {
|
|
return
|
|
} else if k < pivotIndex {
|
|
right = pivotIndex - 1
|
|
} else {
|
|
left = pivotIndex + 1
|
|
}
|
|
}
|
|
}
|
|
|
|
// partition Partition values into less than s[pivot] and greater than s[pivot]
|
|
func partition[S ~[]E, E cmp.Ordered](s S, pivot, left, right int) int {
|
|
s[pivot], s[right] = s[right], s[pivot] // Move pivot to end
|
|
storeIndex := left
|
|
for i := left; i < right; i++ {
|
|
if s[i] < s[right] { // Compare to pivot value
|
|
s[storeIndex], s[i] = s[i], s[storeIndex]
|
|
storeIndex++
|
|
}
|
|
}
|
|
|
|
s[right], s[storeIndex] = s[storeIndex], s[right] // Return pivot Index to position
|
|
return storeIndex
|
|
}
|
|
|
|
// NthElementSliceFunc QuickSelect implementation
|
|
// k is the desired index value, where array[k] is the k+1 smallest element
|
|
// values between s[0, k-1] are guaranteed <= to s[k]
|
|
// values between s[k+1, len(s)] are guaranteed >= to s[k]
|
|
func NthElementSliceFunc[S ~[]E, E any](s S, cmp func(a, b E) int, k int) {
|
|
left := 0
|
|
right := len(s) - 1
|
|
for {
|
|
if left == right {
|
|
return
|
|
}
|
|
pivotIndex := (left + right) >> 1 // Use middle point as pivot. This could probably use random index
|
|
pivotIndex = partitionFunc(s, cmp, pivotIndex, left, right)
|
|
if k == pivotIndex {
|
|
return
|
|
} else if k < pivotIndex {
|
|
right = pivotIndex - 1
|
|
} else {
|
|
left = pivotIndex + 1
|
|
}
|
|
}
|
|
}
|
|
|
|
// partitionFunc Partition values into less than s[pivot] and greater than s[pivot]
|
|
func partitionFunc[S ~[]E, E any](s S, cmp func(a, b E) int, pivot, left, right int) int {
|
|
s[pivot], s[right] = s[right], s[pivot] // Move pivot to end
|
|
storeIndex := left
|
|
for i := left; i < right; i++ {
|
|
if cmp(s[i], s[right]) < 0 { // Compare to pivot value
|
|
s[storeIndex], s[i] = s[i], s[storeIndex]
|
|
storeIndex++
|
|
}
|
|
}
|
|
|
|
s[right], s[storeIndex] = s[storeIndex], s[right] // Return pivot Index to position
|
|
return storeIndex
|
|
}
|