github.com/jingcheng-WU/gonum@v0.9.1-0.20210323123734-f1a2a11a8f7b/stat/sampleuv/withoutreplacement.go (about)

     1  // Copyright ©2017 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 sampleuv
     6  
     7  import (
     8  	"sort"
     9  
    10  	"golang.org/x/exp/rand"
    11  )
    12  
    13  // WithoutReplacement samples len(idxs) integers from [0, n) without replacement.
    14  // That is, upon return the elements of idxs will be unique integers. If source
    15  // is non-nil it will be used to generate random numbers, otherwise the default
    16  // source from the math/rand package will be used.
    17  //
    18  // WithoutReplacement will panic if len(idxs) > n.
    19  func WithoutReplacement(idxs []int, n int, src rand.Source) {
    20  	if len(idxs) == 0 {
    21  		panic("withoutreplacement: zero length input")
    22  	}
    23  	if len(idxs) > n {
    24  		panic("withoutreplacement: impossible size inputs")
    25  	}
    26  
    27  	// There are two algorithms. One is to generate a random permutation
    28  	// and take the first len(idxs) elements. The second is to generate
    29  	// individual random numbers for each element and check uniqueness. The first
    30  	// method scales as O(n), and the second scales as O(len(idxs)^2). Choose
    31  	// the algorithm accordingly.
    32  	if n < len(idxs)*len(idxs) {
    33  		var perm []int
    34  		if src != nil {
    35  			perm = rand.New(src).Perm(n)
    36  		} else {
    37  			perm = rand.Perm(n)
    38  		}
    39  		copy(idxs, perm)
    40  		return
    41  	}
    42  
    43  	// Instead, generate the random numbers directly.
    44  	sorted := make([]int, 0, len(idxs))
    45  	for i := range idxs {
    46  		var r int
    47  		if src != nil {
    48  			r = rand.New(src).Intn(n - i)
    49  		} else {
    50  			r = rand.Intn(n - i)
    51  		}
    52  		for _, v := range sorted {
    53  			if r >= v {
    54  				r++
    55  			}
    56  		}
    57  		idxs[i] = r
    58  		sorted = append(sorted, r)
    59  		sort.Ints(sorted)
    60  	}
    61  }