| 22 | namespace random { |
| 23 | |
| 24 | DistributionSampler::DistributionSampler( |
| 25 | const gtl::ArraySlice<float>& weights) { |
| 26 | DCHECK(!weights.empty()); |
| 27 | int n = weights.size(); |
| 28 | num_ = n; |
| 29 | data_.reset(new std::pair<float, int>[n]); |
| 30 | |
| 31 | std::unique_ptr<double[]> pr(new double[n]); |
| 32 | |
| 33 | double sum = 0.0; |
| 34 | for (int i = 0; i < n; i++) { |
| 35 | sum += weights[i]; |
| 36 | set_alt(i, -1); |
| 37 | } |
| 38 | |
| 39 | // These are long/short items - called high/low because of reserved keywords. |
| 40 | std::vector<int> high; |
| 41 | high.reserve(n); |
| 42 | std::vector<int> low; |
| 43 | low.reserve(n); |
| 44 | |
| 45 | // compute proportional weights |
| 46 | for (int i = 0; i < n; i++) { |
| 47 | double p = (weights[i] * n) / sum; |
| 48 | pr[i] = p; |
| 49 | if (p < 1.0) { |
| 50 | low.push_back(i); |
| 51 | } else { |
| 52 | high.push_back(i); |
| 53 | } |
| 54 | } |
| 55 | |
| 56 | // Now pair high with low. |
| 57 | while (!high.empty() && !low.empty()) { |
| 58 | int l = low.back(); |
| 59 | low.pop_back(); |
| 60 | int h = high.back(); |
| 61 | high.pop_back(); |
| 62 | |
| 63 | set_alt(l, h); |
| 64 | DCHECK_GE(pr[h], 1.0); |
| 65 | double remaining = pr[h] - (1.0 - pr[l]); |
| 66 | pr[h] = remaining; |
| 67 | |
| 68 | if (remaining < 1.0) { |
| 69 | low.push_back(h); |
| 70 | } else { |
| 71 | high.push_back(h); |
| 72 | } |
| 73 | } |
| 74 | // Transfer pr to prob with rounding errors. |
| 75 | for (int i = 0; i < n; i++) { |
| 76 | set_prob(i, pr[i]); |
| 77 | } |
| 78 | // Because of rounding errors, both high and low may have elements, that are |
| 79 | // close to 1.0 prob. |
| 80 | for (size_t i = 0; i < high.size(); i++) { |
| 81 | int idx = high[i]; |