| 180 | } |
| 181 | |
| 182 | void WeightedPicker::Resize(int new_size) { |
| 183 | CHECK_GE(new_size, 0); |
| 184 | if (new_size <= LevelSize(num_levels_ - 1)) { |
| 185 | // The new picker fits in the existing levels. |
| 186 | |
| 187 | // First zero out any of the weights that are being dropped so |
| 188 | // that the levels are correct (only needed when shrinking) |
| 189 | for (int i = new_size; i < N_; i++) { |
| 190 | set_weight(i, 0); |
| 191 | } |
| 192 | |
| 193 | // We do not need to set any new weights when enlarging because |
| 194 | // the unneeded entries always have weight zero. |
| 195 | N_ = new_size; |
| 196 | return; |
| 197 | } |
| 198 | |
| 199 | // We follow the simple strategy of just copying the old |
| 200 | // WeightedPicker into a new WeightedPicker. The cost is |
| 201 | // O(N) regardless. |
| 202 | assert(new_size > N_); |
| 203 | WeightedPicker new_picker(new_size); |
| 204 | int32* dst = new_picker.level_[new_picker.num_levels_ - 1]; |
| 205 | int32* src = this->level_[this->num_levels_ - 1]; |
| 206 | memcpy(dst, src, sizeof(dst[0]) * N_); |
| 207 | memset(dst + N_, 0, sizeof(dst[0]) * (new_size - N_)); |
| 208 | new_picker.RebuildTreeWeights(); |
| 209 | |
| 210 | // Now swap the two pickers |
| 211 | std::swap(new_picker.N_, this->N_); |
| 212 | std::swap(new_picker.num_levels_, this->num_levels_); |
| 213 | std::swap(new_picker.level_, this->level_); |
| 214 | assert(this->N_ == new_size); |
| 215 | } |
| 216 | |
| 217 | } // namespace random |
| 218 | } // namespace tensorflow |