github.com/MetalBlockchain/metalgo@v1.11.9/utils/sampler/uniform_resample.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package sampler
     5  
     6  // uniformResample allows for sampling over a uniform distribution without
     7  // replacement.
     8  //
     9  // Sampling is performed by sampling with replacement and resampling if a
    10  // duplicate is sampled.
    11  //
    12  // Initialization takes O(1) time.
    13  //
    14  // Sampling is performed in O(count) time and O(count) space.
    15  type uniformResample struct {
    16  	rng    *rng
    17  	length uint64
    18  	drawn  map[uint64]struct{}
    19  }
    20  
    21  func (s *uniformResample) Initialize(length uint64) {
    22  	s.length = length
    23  	s.drawn = make(map[uint64]struct{})
    24  }
    25  
    26  func (s *uniformResample) Sample(count int) ([]uint64, bool) {
    27  	s.Reset()
    28  
    29  	results := make([]uint64, count)
    30  	for i := 0; i < count; i++ {
    31  		ret, hasNext := s.Next()
    32  		if !hasNext {
    33  			return nil, false
    34  		}
    35  		results[i] = ret
    36  	}
    37  	return results, true
    38  }
    39  
    40  func (s *uniformResample) Reset() {
    41  	clear(s.drawn)
    42  }
    43  
    44  func (s *uniformResample) Next() (uint64, bool) {
    45  	i := uint64(len(s.drawn))
    46  	if i >= s.length {
    47  		return 0, false
    48  	}
    49  
    50  	for {
    51  		draw := s.rng.Uint64Inclusive(s.length - 1)
    52  		if _, ok := s.drawn[draw]; ok {
    53  			continue
    54  		}
    55  		s.drawn[draw] = struct{}{}
    56  		return draw, true
    57  	}
    58  }