/
medians.go
106 lines (96 loc) · 2.34 KB
/
medians.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package kdtree
import (
"sort"
"golang.org/x/exp/rand"
)
// Partition partitions list such that all elements less than the value at
// pivot prior to the call are placed before that element and all elements
// greater than that value are placed after it. The final location of the
// element at pivot prior to the call is returned.
func Partition(list sort.Interface, pivot int) int {
var index, last int
if last = list.Len() - 1; last < 0 {
return -1
}
list.Swap(pivot, last)
for i := 0; i < last; i++ {
if !list.Less(last, i) {
list.Swap(index, i)
index++
}
}
list.Swap(last, index)
return index
}
// SortSlicer satisfies the sort.Interface and is able to slice itself.
type SortSlicer interface {
sort.Interface
Slice(start, end int) SortSlicer
}
// Select partitions list such that all elements less than the kth element
// are placed before k in the resulting list and all elements greater than
// it are placed after the position k.
func Select(list SortSlicer, k int) int {
var (
start int
end = list.Len()
)
if k >= end {
if k == 0 {
return 0
}
panic("kdtree: index out of range")
}
if start == end-1 {
return k
}
for {
if start == end {
panic("kdtree: internal inconsistency")
}
sub := list.Slice(start, end)
pivot := Partition(sub, rand.Intn(sub.Len()))
switch {
case pivot == k:
return k
case k < pivot:
end = pivot + start
default:
k -= pivot
start += pivot
}
}
}
func min(a, b int) int {
if a < b {
return a
}
return b
}
// MedianOfMedians returns the index to the median value of the medians
// of groups of 5 consecutive elements.
func MedianOfMedians(list SortSlicer) int {
n := list.Len() / 5
for i := 0; i < n; i++ {
left := i * 5
sub := list.Slice(left, min(left+5, list.Len()-1))
Select(sub, 2)
list.Swap(i, left+2)
}
Select(list.Slice(0, min(n, list.Len()-1)), min(list.Len(), n/2))
return n / 2
}
// MedianOfRandoms returns the index to the median value of up to n randomly
// chosen elements in list.
func MedianOfRandoms(list SortSlicer, n int) int {
if l := list.Len(); l < n {
n = l
} else {
rand.Shuffle(n, func(i, j int) { list.Swap(i, j) })
}
Select(list.Slice(0, n), n/2)
return n / 2
}