1// Copyright 2018 Developers of the Rand project.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9use super::{Error, Weight};
10use crate::distr::uniform::{SampleBorrow, SampleUniform, UniformSampler};
11use crate::distr::Distribution;
12use crate::Rng;
13
14// Note that this whole module is only imported if feature="alloc" is enabled.
15use alloc::vec::Vec;
16use core::fmt::{self, Debug};
17
18#[cfg(feature = "serde")]
19use serde::{Deserialize, Serialize};
20
21/// A distribution using weighted sampling of discrete items.
22///
23/// Sampling a `WeightedIndex` distribution returns the index of a randomly
24/// selected element from the iterator used when the `WeightedIndex` was
25/// created. The chance of a given element being picked is proportional to the
26/// weight of the element. The weights can use any type `X` for which an
27/// implementation of [`Uniform<X>`] exists. The implementation guarantees that
28/// elements with zero weight are never picked, even when the weights are
29/// floating point numbers.
30///
31/// # Performance
32///
33/// Time complexity of sampling from `WeightedIndex` is `O(log N)` where
34/// `N` is the number of weights.
35/// See also [`rand_distr::weighted`] for alternative implementations supporting
36/// potentially-faster sampling or a more easily modifiable tree structure.
37///
38/// A `WeightedIndex<X>` contains a `Vec<X>` and a [`Uniform<X>`] and so its
39/// size is the sum of the size of those objects, possibly plus some alignment.
40///
41/// Creating a `WeightedIndex<X>` will allocate enough space to hold `N - 1`
42/// weights of type `X`, where `N` is the number of weights. However, since
43/// `Vec` doesn't guarantee a particular growth strategy, additional memory
44/// might be allocated but not used. Since the `WeightedIndex` object also
45/// contains an instance of `X::Sampler`, this might cause additional allocations,
46/// though for primitive types, [`Uniform<X>`] doesn't allocate any memory.
47///
48/// Sampling from `WeightedIndex` will result in a single call to
49/// `Uniform<X>::sample` (method of the [`Distribution`] trait), which typically
50/// will request a single value from the underlying [`RngCore`], though the
51/// exact number depends on the implementation of `Uniform<X>::sample`.
52///
53/// # Example
54///
55/// ```
56/// use rand::prelude::*;
57/// use rand::distr::weighted::WeightedIndex;
58///
59/// let choices = ['a', 'b', 'c'];
60/// let weights = [2, 1, 1];
61/// let dist = WeightedIndex::new(&weights).unwrap();
62/// let mut rng = rand::rng();
63/// for _ in 0..100 {
64/// // 50% chance to print 'a', 25% chance to print 'b', 25% chance to print 'c'
65/// println!("{}", choices[dist.sample(&mut rng)]);
66/// }
67///
68/// let items = [('a', 0.0), ('b', 3.0), ('c', 7.0)];
69/// let dist2 = WeightedIndex::new(items.iter().map(|item| item.1)).unwrap();
70/// for _ in 0..100 {
71/// // 0% chance to print 'a', 30% chance to print 'b', 70% chance to print 'c'
72/// println!("{}", items[dist2.sample(&mut rng)].0);
73/// }
74/// ```
75///
76/// [`Uniform<X>`]: crate::distr::Uniform
77/// [`RngCore`]: crate::RngCore
78/// [`rand_distr::weighted`]: https://docs.rs/rand_distr/latest/rand_distr/weighted/index.html
79#[derive(Debug, Clone, PartialEq)]
80#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
81pub struct WeightedIndex<X: SampleUniform + PartialOrd> {
82 cumulative_weights: Vec<X>,
83 total_weight: X,
84 weight_distribution: X::Sampler,
85}
86
87impl<X: SampleUniform + PartialOrd> WeightedIndex<X> {
88 /// Creates a new a `WeightedIndex` [`Distribution`] using the values
89 /// in `weights`. The weights can use any type `X` for which an
90 /// implementation of [`Uniform<X>`] exists.
91 ///
92 /// Error cases:
93 /// - [`Error::InvalidInput`] when the iterator `weights` is empty.
94 /// - [`Error::InvalidWeight`] when a weight is not-a-number or negative.
95 /// - [`Error::InsufficientNonZero`] when the sum of all weights is zero.
96 /// - [`Error::Overflow`] when the sum of all weights overflows.
97 ///
98 /// [`Uniform<X>`]: crate::distr::uniform::Uniform
99 pub fn new<I>(weights: I) -> Result<WeightedIndex<X>, Error>
100 where
101 I: IntoIterator,
102 I::Item: SampleBorrow<X>,
103 X: Weight,
104 {
105 let mut iter = weights.into_iter();
106 let mut total_weight: X = iter.next().ok_or(Error::InvalidInput)?.borrow().clone();
107
108 let zero = X::ZERO;
109 if !(total_weight >= zero) {
110 return Err(Error::InvalidWeight);
111 }
112
113 let mut weights = Vec::<X>::with_capacity(iter.size_hint().0);
114 for w in iter {
115 // Note that `!(w >= x)` is not equivalent to `w < x` for partially
116 // ordered types due to NaNs which are equal to nothing.
117 if !(w.borrow() >= &zero) {
118 return Err(Error::InvalidWeight);
119 }
120 weights.push(total_weight.clone());
121
122 if let Err(()) = total_weight.checked_add_assign(w.borrow()) {
123 return Err(Error::Overflow);
124 }
125 }
126
127 if total_weight == zero {
128 return Err(Error::InsufficientNonZero);
129 }
130 let distr = X::Sampler::new(zero, total_weight.clone()).unwrap();
131
132 Ok(WeightedIndex {
133 cumulative_weights: weights,
134 total_weight,
135 weight_distribution: distr,
136 })
137 }
138
139 /// Update a subset of weights, without changing the number of weights.
140 ///
141 /// `new_weights` must be sorted by the index.
142 ///
143 /// Using this method instead of `new` might be more efficient if only a small number of
144 /// weights is modified. No allocations are performed, unless the weight type `X` uses
145 /// allocation internally.
146 ///
147 /// In case of error, `self` is not modified. Error cases:
148 /// - [`Error::InvalidInput`] when `new_weights` are not ordered by
149 /// index or an index is too large.
150 /// - [`Error::InvalidWeight`] when a weight is not-a-number or negative.
151 /// - [`Error::InsufficientNonZero`] when the sum of all weights is zero.
152 /// Note that due to floating-point loss of precision, this case is not
153 /// always correctly detected; usage of a fixed-point weight type may be
154 /// preferred.
155 ///
156 /// Updates take `O(N)` time. If you need to frequently update weights, consider
157 /// [`rand_distr::weighted_tree`](https://docs.rs/rand_distr/*/rand_distr/weighted_tree/index.html)
158 /// as an alternative where an update is `O(log N)`.
159 pub fn update_weights(&mut self, new_weights: &[(usize, &X)]) -> Result<(), Error>
160 where
161 X: for<'a> core::ops::AddAssign<&'a X>
162 + for<'a> core::ops::SubAssign<&'a X>
163 + Clone
164 + Default,
165 {
166 if new_weights.is_empty() {
167 return Ok(());
168 }
169
170 let zero = <X as Default>::default();
171
172 let mut total_weight = self.total_weight.clone();
173
174 // Check for errors first, so we don't modify `self` in case something
175 // goes wrong.
176 let mut prev_i = None;
177 for &(i, w) in new_weights {
178 if let Some(old_i) = prev_i {
179 if old_i >= i {
180 return Err(Error::InvalidInput);
181 }
182 }
183 if !(*w >= zero) {
184 return Err(Error::InvalidWeight);
185 }
186 if i > self.cumulative_weights.len() {
187 return Err(Error::InvalidInput);
188 }
189
190 let mut old_w = if i < self.cumulative_weights.len() {
191 self.cumulative_weights[i].clone()
192 } else {
193 self.total_weight.clone()
194 };
195 if i > 0 {
196 old_w -= &self.cumulative_weights[i - 1];
197 }
198
199 total_weight -= &old_w;
200 total_weight += w;
201 prev_i = Some(i);
202 }
203 if total_weight <= zero {
204 return Err(Error::InsufficientNonZero);
205 }
206
207 // Update the weights. Because we checked all the preconditions in the
208 // previous loop, this should never panic.
209 let mut iter = new_weights.iter();
210
211 let mut prev_weight = zero.clone();
212 let mut next_new_weight = iter.next();
213 let &(first_new_index, _) = next_new_weight.unwrap();
214 let mut cumulative_weight = if first_new_index > 0 {
215 self.cumulative_weights[first_new_index - 1].clone()
216 } else {
217 zero.clone()
218 };
219 for i in first_new_index..self.cumulative_weights.len() {
220 match next_new_weight {
221 Some(&(j, w)) if i == j => {
222 cumulative_weight += w;
223 next_new_weight = iter.next();
224 }
225 _ => {
226 let mut tmp = self.cumulative_weights[i].clone();
227 tmp -= &prev_weight; // We know this is positive.
228 cumulative_weight += &tmp;
229 }
230 }
231 prev_weight = cumulative_weight.clone();
232 core::mem::swap(&mut prev_weight, &mut self.cumulative_weights[i]);
233 }
234
235 self.total_weight = total_weight;
236 self.weight_distribution = X::Sampler::new(zero, self.total_weight.clone()).unwrap();
237
238 Ok(())
239 }
240}
241
242/// A lazy-loading iterator over the weights of a `WeightedIndex` distribution.
243/// This is returned by [`WeightedIndex::weights`].
244pub struct WeightedIndexIter<'a, X: SampleUniform + PartialOrd> {
245 weighted_index: &'a WeightedIndex<X>,
246 index: usize,
247}
248
249impl<X> Debug for WeightedIndexIter<'_, X>
250where
251 X: SampleUniform + PartialOrd + Debug,
252 X::Sampler: Debug,
253{
254 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
255 f&mut DebugStruct<'_, '_>.debug_struct("WeightedIndexIter")
256 .field("weighted_index", &self.weighted_index)
257 .field(name:"index", &self.index)
258 .finish()
259 }
260}
261
262impl<X> Clone for WeightedIndexIter<'_, X>
263where
264 X: SampleUniform + PartialOrd,
265{
266 fn clone(&self) -> Self {
267 WeightedIndexIter {
268 weighted_index: self.weighted_index,
269 index: self.index,
270 }
271 }
272}
273
274impl<X> Iterator for WeightedIndexIter<'_, X>
275where
276 X: for<'b> dyncore::ops::SubAssign<&'b X> + SampleUniform + PartialOrd + Clone,
277{
278 type Item = X;
279
280 fn next(&mut self) -> Option<Self::Item> {
281 match self.weighted_index.weight(self.index) {
282 None => None,
283 Some(weight: X) => {
284 self.index += 1;
285 Some(weight)
286 }
287 }
288 }
289}
290
291impl<X: SampleUniform + PartialOrd + Clone> WeightedIndex<X> {
292 /// Returns the weight at the given index, if it exists.
293 ///
294 /// If the index is out of bounds, this will return `None`.
295 ///
296 /// # Example
297 ///
298 /// ```
299 /// use rand::distr::weighted::WeightedIndex;
300 ///
301 /// let weights = [0, 1, 2];
302 /// let dist = WeightedIndex::new(&weights).unwrap();
303 /// assert_eq!(dist.weight(0), Some(0));
304 /// assert_eq!(dist.weight(1), Some(1));
305 /// assert_eq!(dist.weight(2), Some(2));
306 /// assert_eq!(dist.weight(3), None);
307 /// ```
308 pub fn weight(&self, index: usize) -> Option<X>
309 where
310 X: for<'a> core::ops::SubAssign<&'a X>,
311 {
312 use core::cmp::Ordering::*;
313
314 let mut weight = match index.cmp(&self.cumulative_weights.len()) {
315 Less => self.cumulative_weights[index].clone(),
316 Equal => self.total_weight.clone(),
317 Greater => return None,
318 };
319
320 if index > 0 {
321 weight -= &self.cumulative_weights[index - 1];
322 }
323 Some(weight)
324 }
325
326 /// Returns a lazy-loading iterator containing the current weights of this distribution.
327 ///
328 /// If this distribution has not been updated since its creation, this will return the
329 /// same weights as were passed to `new`.
330 ///
331 /// # Example
332 ///
333 /// ```
334 /// use rand::distr::weighted::WeightedIndex;
335 ///
336 /// let weights = [1, 2, 3];
337 /// let mut dist = WeightedIndex::new(&weights).unwrap();
338 /// assert_eq!(dist.weights().collect::<Vec<_>>(), vec![1, 2, 3]);
339 /// dist.update_weights(&[(0, &2)]).unwrap();
340 /// assert_eq!(dist.weights().collect::<Vec<_>>(), vec![2, 2, 3]);
341 /// ```
342 pub fn weights(&self) -> WeightedIndexIter<'_, X>
343 where
344 X: for<'a> core::ops::SubAssign<&'a X>,
345 {
346 WeightedIndexIter {
347 weighted_index: self,
348 index: 0,
349 }
350 }
351
352 /// Returns the sum of all weights in this distribution.
353 pub fn total_weight(&self) -> X {
354 self.total_weight.clone()
355 }
356}
357
358impl<X> Distribution<usize> for WeightedIndex<X>
359where
360 X: SampleUniform + PartialOrd,
361{
362 fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
363 let chosen_weight: X = self.weight_distribution.sample(rng);
364 // Find the first item which has a weight *higher* than the chosen weight.
365 self.cumulative_weights
366 .partition_point(|w: &X| w <= &chosen_weight)
367 }
368}
369
370#[cfg(test)]
371mod test {
372 use super::*;
373
374 #[cfg(feature = "serde")]
375 #[test]
376 fn test_weightedindex_serde() {
377 let weighted_index = WeightedIndex::new([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).unwrap();
378
379 let ser_weighted_index = bincode::serialize(&weighted_index).unwrap();
380 let de_weighted_index: WeightedIndex<i32> =
381 bincode::deserialize(&ser_weighted_index).unwrap();
382
383 assert_eq!(
384 de_weighted_index.cumulative_weights,
385 weighted_index.cumulative_weights
386 );
387 assert_eq!(de_weighted_index.total_weight, weighted_index.total_weight);
388 }
389
390 #[test]
391 fn test_accepting_nan() {
392 assert_eq!(
393 WeightedIndex::new([f32::NAN, 0.5]).unwrap_err(),
394 Error::InvalidWeight,
395 );
396 assert_eq!(
397 WeightedIndex::new([f32::NAN]).unwrap_err(),
398 Error::InvalidWeight,
399 );
400 assert_eq!(
401 WeightedIndex::new([0.5, f32::NAN]).unwrap_err(),
402 Error::InvalidWeight,
403 );
404
405 assert_eq!(
406 WeightedIndex::new([0.5, 7.0])
407 .unwrap()
408 .update_weights(&[(0, &f32::NAN)])
409 .unwrap_err(),
410 Error::InvalidWeight,
411 )
412 }
413
414 #[test]
415 #[cfg_attr(miri, ignore)] // Miri is too slow
416 fn test_weightedindex() {
417 let mut r = crate::test::rng(700);
418 const N_REPS: u32 = 5000;
419 let weights = [1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7];
420 let total_weight = weights.iter().sum::<u32>() as f32;
421
422 let verify = |result: [i32; 14]| {
423 for (i, count) in result.iter().enumerate() {
424 let exp = (weights[i] * N_REPS) as f32 / total_weight;
425 let mut err = (*count as f32 - exp).abs();
426 if err != 0.0 {
427 err /= exp;
428 }
429 assert!(err <= 0.25);
430 }
431 };
432
433 // WeightedIndex from vec
434 let mut chosen = [0i32; 14];
435 let distr = WeightedIndex::new(weights.to_vec()).unwrap();
436 for _ in 0..N_REPS {
437 chosen[distr.sample(&mut r)] += 1;
438 }
439 verify(chosen);
440
441 // WeightedIndex from slice
442 chosen = [0i32; 14];
443 let distr = WeightedIndex::new(&weights[..]).unwrap();
444 for _ in 0..N_REPS {
445 chosen[distr.sample(&mut r)] += 1;
446 }
447 verify(chosen);
448
449 // WeightedIndex from iterator
450 chosen = [0i32; 14];
451 let distr = WeightedIndex::new(weights.iter()).unwrap();
452 for _ in 0..N_REPS {
453 chosen[distr.sample(&mut r)] += 1;
454 }
455 verify(chosen);
456
457 for _ in 0..5 {
458 assert_eq!(WeightedIndex::new([0, 1]).unwrap().sample(&mut r), 1);
459 assert_eq!(WeightedIndex::new([1, 0]).unwrap().sample(&mut r), 0);
460 assert_eq!(
461 WeightedIndex::new([0, 0, 0, 0, 10, 0])
462 .unwrap()
463 .sample(&mut r),
464 4
465 );
466 }
467
468 assert_eq!(
469 WeightedIndex::new(&[10][0..0]).unwrap_err(),
470 Error::InvalidInput
471 );
472 assert_eq!(
473 WeightedIndex::new([0]).unwrap_err(),
474 Error::InsufficientNonZero
475 );
476 assert_eq!(
477 WeightedIndex::new([10, 20, -1, 30]).unwrap_err(),
478 Error::InvalidWeight
479 );
480 assert_eq!(
481 WeightedIndex::new([-10, 20, 1, 30]).unwrap_err(),
482 Error::InvalidWeight
483 );
484 assert_eq!(WeightedIndex::new([-10]).unwrap_err(), Error::InvalidWeight);
485 }
486
487 #[test]
488 fn test_update_weights() {
489 let data = [
490 (
491 &[10u32, 2, 3, 4][..],
492 &[(1, &100), (2, &4)][..], // positive change
493 &[10, 100, 4, 4][..],
494 ),
495 (
496 &[1u32, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7][..],
497 &[(2, &1), (5, &1), (13, &100)][..], // negative change and last element
498 &[1u32, 2, 1, 0, 5, 1, 7, 1, 2, 3, 4, 5, 6, 100][..],
499 ),
500 ];
501
502 for (weights, update, expected_weights) in data.iter() {
503 let total_weight = weights.iter().sum::<u32>();
504 let mut distr = WeightedIndex::new(weights.to_vec()).unwrap();
505 assert_eq!(distr.total_weight, total_weight);
506
507 distr.update_weights(update).unwrap();
508 let expected_total_weight = expected_weights.iter().sum::<u32>();
509 let expected_distr = WeightedIndex::new(expected_weights.to_vec()).unwrap();
510 assert_eq!(distr.total_weight, expected_total_weight);
511 assert_eq!(distr.total_weight, expected_distr.total_weight);
512 assert_eq!(distr.cumulative_weights, expected_distr.cumulative_weights);
513 }
514 }
515
516 #[test]
517 fn test_update_weights_errors() {
518 let data = [
519 (
520 &[1i32, 0, 0][..],
521 &[(0, &0)][..],
522 Error::InsufficientNonZero,
523 ),
524 (
525 &[10, 10, 10, 10][..],
526 &[(1, &-11)][..],
527 Error::InvalidWeight, // A weight is negative
528 ),
529 (
530 &[1, 2, 3, 4, 5][..],
531 &[(1, &5), (0, &5)][..], // Wrong order
532 Error::InvalidInput,
533 ),
534 (
535 &[1][..],
536 &[(1, &1)][..], // Index too large
537 Error::InvalidInput,
538 ),
539 ];
540
541 for (weights, update, err) in data.iter() {
542 let total_weight = weights.iter().sum::<i32>();
543 let mut distr = WeightedIndex::new(weights.to_vec()).unwrap();
544 assert_eq!(distr.total_weight, total_weight);
545 match distr.update_weights(update) {
546 Ok(_) => panic!("Expected update_weights to fail, but it succeeded"),
547 Err(e) => assert_eq!(e, *err),
548 }
549 }
550 }
551
552 #[test]
553 fn test_weight_at() {
554 let data = [
555 &[1][..],
556 &[10, 2, 3, 4][..],
557 &[1, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7][..],
558 &[u32::MAX][..],
559 ];
560
561 for weights in data.iter() {
562 let distr = WeightedIndex::new(weights.to_vec()).unwrap();
563 for (i, weight) in weights.iter().enumerate() {
564 assert_eq!(distr.weight(i), Some(*weight));
565 }
566 assert_eq!(distr.weight(weights.len()), None);
567 }
568 }
569
570 #[test]
571 fn test_weights() {
572 let data = [
573 &[1][..],
574 &[10, 2, 3, 4][..],
575 &[1, 2, 3, 0, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7][..],
576 &[u32::MAX][..],
577 ];
578
579 for weights in data.iter() {
580 let distr = WeightedIndex::new(weights.to_vec()).unwrap();
581 assert_eq!(distr.weights().collect::<Vec<_>>(), weights.to_vec());
582 }
583 }
584
585 #[test]
586 fn value_stability() {
587 fn test_samples<X: Weight + SampleUniform + PartialOrd, I>(
588 weights: I,
589 buf: &mut [usize],
590 expected: &[usize],
591 ) where
592 I: IntoIterator,
593 I::Item: SampleBorrow<X>,
594 {
595 assert_eq!(buf.len(), expected.len());
596 let distr = WeightedIndex::new(weights).unwrap();
597 let mut rng = crate::test::rng(701);
598 for r in buf.iter_mut() {
599 *r = rng.sample(&distr);
600 }
601 assert_eq!(buf, expected);
602 }
603
604 let mut buf = [0; 10];
605 test_samples(
606 [1i32, 1, 1, 1, 1, 1, 1, 1, 1],
607 &mut buf,
608 &[0, 6, 2, 6, 3, 4, 7, 8, 2, 5],
609 );
610 test_samples(
611 [0.7f32, 0.1, 0.1, 0.1],
612 &mut buf,
613 &[0, 0, 0, 1, 0, 0, 2, 3, 0, 0],
614 );
615 test_samples(
616 [1.0f64, 0.999, 0.998, 0.997],
617 &mut buf,
618 &[2, 2, 1, 3, 2, 1, 3, 3, 2, 1],
619 );
620 }
621
622 #[test]
623 fn weighted_index_distributions_can_be_compared() {
624 assert_eq!(WeightedIndex::new([1, 2]), WeightedIndex::new([1, 2]));
625 }
626
627 #[test]
628 fn overflow() {
629 assert_eq!(WeightedIndex::new([2, usize::MAX]), Err(Error::Overflow));
630 }
631}
632