github.com/niubaoshu/goutils@v0.0.0-20180828035119-e8e576f66c2b/rand/mRandN.go (about)

     1  // 该程序实现在m个int中选取n个
     2  package rand
     3  
     4  import (
     5  	"math/rand"
     6  	"strconv"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/niubaoshu/goutils"
    11  )
    12  
    13  type MrandN struct {
    14  	lock   sync.Mutex
    15  	bucket []int
    16  	bktLen int
    17  	rand   *rand.Rand
    18  }
    19  
    20  const defaultLength = 1024
    21  
    22  var defaultRand = NewMrandNWithLen(defaultLength)
    23  
    24  func NewMrandN(bkt []int) *MrandN {
    25  	return &MrandN{bktLen: len(bkt), bucket: bkt, rand: rand.New(rand.NewSource(time.Now().UnixNano()))}
    26  }
    27  
    28  func NewMrandNWithLen(l int) *MrandN {
    29  	bkt := make([]int, l)
    30  	for i := 0; i < l; i++ {
    31  		bkt[i] = i
    32  	}
    33  	return &MrandN{bktLen: l, bucket: bkt, rand: rand.New(rand.NewSource(time.Now().UnixNano()))}
    34  }
    35  
    36  func (mrn *MrandN) SelectN(n int, data []int) []int {
    37  	if mrn.bktLen < n {
    38  		panic("bucket length " + strconv.Itoa(mrn.bktLen) + " not enough " + strconv.Itoa(n))
    39  	}
    40  
    41  	bkt := mrn.bucket
    42  	bl := mrn.bktLen
    43  	r := mrn.rand
    44  	last := mrn.bktLen - 1
    45  
    46  	buf, l := goutils.EnlargeInt(data, n)
    47  	temp := buf[l:]
    48  
    49  	mrn.lock.Lock()
    50  	for i := 0; i < n; i++ {
    51  		idx := r.Intn(bl - i)
    52  		temp[i] = bkt[idx]
    53  		bkt[idx], bkt[last-i] = bkt[last-i], bkt[idx]
    54  	}
    55  	mrn.lock.Unlock()
    56  	return buf[:l+n]
    57  }
    58  
    59  func (mrn *MrandN) Perm(data []int) []int {
    60  	return mrn.SelectN(mrn.bktLen, data)
    61  }
    62  
    63  func MSelectN(m, n int, data []int) []int {
    64  	var r = defaultRand
    65  	if m > defaultLength {
    66  		r = NewMrandNWithLen(m)
    67  	}
    68  	return r.SelectN(n, data)
    69  }
    70  
    71  func Perm(n int, data []int) []int {
    72  	return MSelectN(n, n, data)
    73  }