You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
87 lines
2.2 KiB
87 lines
2.2 KiB
package autopilot |
|
|
|
import ( |
|
"errors" |
|
"fmt" |
|
"math/rand" |
|
) |
|
|
|
// ErrNoPositive is returned from weightedChoice when there are no positive |
|
// weights left to choose from. |
|
var ErrNoPositive = errors.New("no positive weights left") |
|
|
|
// weightedChoice draws a random index from the slice of weights, with a |
|
// probability propotional to the weight at the given index. |
|
func weightedChoice(w []float64) (int, error) { |
|
// Calculate the sum of weights. |
|
var sum float64 |
|
for _, v := range w { |
|
sum += v |
|
} |
|
|
|
if sum <= 0 { |
|
return 0, ErrNoPositive |
|
} |
|
|
|
// Pick a random number in the range [0.0, 1.0) and multiply it with |
|
// the sum of weights. Then we'll iterate the weights until the number |
|
// goes below 0. This means that each index is picked with a probablity |
|
// equal to their normalized score. |
|
// |
|
// Example: |
|
// Items with scores [1, 5, 2, 2] |
|
// Normalized scores [0.1, 0.5, 0.2, 0.2] |
|
// Imagine they each occupy a "range" equal to their normalized score |
|
// in [0, 1.0]: |
|
// [|-0.1-||-----0.5-----||--0.2--||--0.2--|] |
|
// The following loop is now equivalent to "hitting" the intervals. |
|
r := rand.Float64() * sum |
|
for i := range w { |
|
r -= w[i] |
|
if r <= 0 { |
|
return i, nil |
|
} |
|
} |
|
|
|
return 0, fmt.Errorf("unable to make choice") |
|
} |
|
|
|
// chooseN picks at random min[n, len(s)] nodes if from the NodeScore map, with |
|
// a probability weighted by their score. |
|
func chooseN(n uint32, s map[NodeID]*NodeScore) ( |
|
map[NodeID]*NodeScore, error) { |
|
|
|
// Keep track of the number of nodes not yet chosen, in addition to |
|
// their scores and NodeIDs. |
|
rem := len(s) |
|
scores := make([]float64, len(s)) |
|
nodeIDs := make([]NodeID, len(s)) |
|
i := 0 |
|
for k, v := range s { |
|
scores[i] = v.Score |
|
nodeIDs[i] = k |
|
i++ |
|
} |
|
|
|
// Pick a weighted choice from the remaining nodes as long as there are |
|
// nodes left, and we haven't already picked n. |
|
chosen := make(map[NodeID]*NodeScore) |
|
for len(chosen) < int(n) && rem > 0 { |
|
choice, err := weightedChoice(scores) |
|
if err == ErrNoPositive { |
|
return chosen, nil |
|
} else if err != nil { |
|
return nil, err |
|
} |
|
|
|
nID := nodeIDs[choice] |
|
|
|
chosen[nID] = s[nID] |
|
|
|
// We set the score of the chosen node to 0, so it won't be |
|
// picked the next iteration. |
|
scores[choice] = 0 |
|
} |
|
|
|
return chosen, nil |
|
}
|
|
|