github.com/gopherd/gonum@v0.0.4/spatial/kdtree/medians.go (about)

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