github.com/decred/dcrlnd@v0.7.6/autopilot/choice.go (about)

     1  package autopilot
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"math/rand"
     7  )
     8  
     9  // ErrNoPositive is returned from weightedChoice when there are no positive
    10  // weights left to choose from.
    11  var ErrNoPositive = errors.New("no positive weights left")
    12  
    13  // weightedChoice draws a random index from the slice of weights, with a
    14  // probability propotional to the weight at the given index.
    15  func weightedChoice(w []float64) (int, error) {
    16  	// Calculate the sum of weights.
    17  	var sum float64
    18  	for _, v := range w {
    19  		sum += v
    20  	}
    21  
    22  	if sum <= 0 {
    23  		return 0, ErrNoPositive
    24  	}
    25  
    26  	// Pick a random number in the range [0.0, 1.0) and multiply it with
    27  	// the sum of weights. Then we'll iterate the weights until the number
    28  	// goes below 0. This means that each index is picked with a probablity
    29  	// equal to their normalized score.
    30  	//
    31  	// Example:
    32  	// Items with scores [1, 5, 2, 2]
    33  	// Normalized scores [0.1, 0.5, 0.2, 0.2]
    34  	// Imagine they each occupy a "range" equal to their normalized score
    35  	// in [0, 1.0]:
    36  	// [|-0.1-||-----0.5-----||--0.2--||--0.2--|]
    37  	// The following loop is now equivalent to "hitting" the intervals.
    38  	r := rand.Float64() * sum
    39  	for i := range w {
    40  		r -= w[i]
    41  		if r <= 0 {
    42  			return i, nil
    43  		}
    44  	}
    45  
    46  	return 0, fmt.Errorf("unable to make choice")
    47  }
    48  
    49  // chooseN picks at random min[n, len(s)] nodes if from the NodeScore map, with
    50  // a probability weighted by their score.
    51  func chooseN(n uint32, s map[NodeID]*NodeScore) (
    52  	map[NodeID]*NodeScore, error) {
    53  
    54  	// Keep track of the number of nodes not yet chosen, in addition to
    55  	// their scores and NodeIDs.
    56  	rem := len(s)
    57  	scores := make([]float64, len(s))
    58  	nodeIDs := make([]NodeID, len(s))
    59  	i := 0
    60  	for k, v := range s {
    61  		scores[i] = v.Score
    62  		nodeIDs[i] = k
    63  		i++
    64  	}
    65  
    66  	// Pick a weighted choice from the remaining nodes as long as there are
    67  	// nodes left, and we haven't already picked n.
    68  	chosen := make(map[NodeID]*NodeScore)
    69  	for len(chosen) < int(n) && rem > 0 {
    70  		choice, err := weightedChoice(scores)
    71  		if err == ErrNoPositive {
    72  			return chosen, nil
    73  		} else if err != nil {
    74  			return nil, err
    75  		}
    76  
    77  		nID := nodeIDs[choice]
    78  
    79  		chosen[nID] = s[nID]
    80  
    81  		// We set the score of the chosen node to 0, so it won't be
    82  		// picked the next iteration.
    83  		scores[choice] = 0
    84  	}
    85  
    86  	return chosen, nil
    87  }