github.com/biogo/store@v0.0.0-20201120204734-aad293a2328f/kdtree/medians.go (about)

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