| 107 | } |
| 108 | |
| 109 | int WeightedPicker::PickAt(int32 weight_index) const { |
| 110 | if (weight_index < 0 || weight_index >= total_weight()) return -1; |
| 111 | |
| 112 | int32 position = weight_index; |
| 113 | int index = 0; |
| 114 | |
| 115 | for (int l = 1; l < num_levels_; l++) { |
| 116 | // Pick left or right child of "level_[l-1][index]" |
| 117 | const int32 left_weight = level_[l][2 * index]; |
| 118 | if (position < left_weight) { |
| 119 | // Descend to left child |
| 120 | index = 2 * index; |
| 121 | } else { |
| 122 | // Descend to right child |
| 123 | index = 2 * index + 1; |
| 124 | position -= left_weight; |
| 125 | } |
| 126 | } |
| 127 | CHECK_GE(index, 0); |
| 128 | CHECK_LT(index, N_); |
| 129 | CHECK_LE(position, level_[num_levels_ - 1][index]); |
| 130 | return index; |
| 131 | } |
| 132 | |
| 133 | void WeightedPicker::set_weight(int index, int32 weight) { |
| 134 | assert(index >= 0); |