weighted_list/weighted_list.rs
1use std::{
2 fmt,
3 iter,
4 ops::*,
5};
6
7use bon::{bon};
8use itertools::Itertools;
9use num_traits as nums;
10use rand::{
11 prelude::*,
12 seq::SliceRandom,
13};
14
15use crate::root::*;
16use crate::WeightedItem;
17
18
19/// A homogeneous list of weighted items with values of type `V` and weights of numerical type `W`.
20///
21/// Near-identical to `Vec<T>`, but stores `WeightedItem<V,W>` objects instead. You can think of it like a `Vec<WeightedItem<V,W>>`.
22///
23/// # Usage
24///
25/// ```
26/// # use weighted_list::*;
27/// let list: WeightedList<String, i32> = wlist![
28/// (2, String::from("sup")),
29/// (3, String::from("nova")),
30/// (5, String::from("shard")),
31/// ];
32///
33/// for item in &list {
34/// println!("{item}");
35/// }
36///
37/// if let Some(result) = list.select_random_value(&mut rand::rng()) {
38/// println!("{}", result);
39/// }
40/// ```
41///
42/// # Tips
43///
44/// - Most methods return `&Self` or `&mut Self`, allowing you to chain methods. Here's a contrived example:
45///
46/// ```
47/// # use weighted_list::*;
48/// let mut list = wlist![(2, "sup")];
49///
50/// list.push_value("sup")
51/// .merge_duplicates()
52/// .prune()
53/// .len();
54/// ```
55#[derive(Debug, Clone, Eq, PartialEq, Hash, Default)]
56pub struct WeightedList<V, W: Weight>
57{
58 data: Vec<WeightedItem<V,W>>
59}
60
61// == CONSTRUCTORS == //
62impl<V, W: Weight> WeightedList<V,W>
63{
64 /// Construct an empty `WeightedList`.
65 pub fn new() -> Self
66 {
67 Self { data: Vec::new() }
68 }
69
70 /// Construct an empty `WeightedList` with the specified capacity.
71 pub fn with_capacity(capacity: usize) -> Self
72 {
73 Self { data: Vec::with_capacity(capacity) }
74 }
75
76 /// Construct a `WeightedList` from an iterable of `(weight, value)` pairs.
77 pub fn init<I>(items: I) -> Self
78 where I: IntoIterator<Item = (W, V)>
79 {
80 Self {
81 data: items.into_iter().map(
82 |(weight, value)|
83 WeightedItem::new(weight, value)
84 ).collect::<Vec<WeightedItem<V,W>>>()
85 }
86 }
87}
88
89/// Construct a `WeightedList` from the provided `(weight, value)`` pairs.
90///
91/// # Usage
92///
93/// ```
94/// # use weighted_list::*;
95/// let list = wlist![
96/// (2, String::from("sup")),
97/// (3, String::from("nova")),
98/// (5, String::from("shard")),
99/// ];
100/// ```
101#[macro_export]
102macro_rules! wlist {
103 ( $( $item: expr ),* $(,)? ) => {
104 WeightedList::init([
105 $( $item, )*
106 ])
107 };
108}
109
110// == ACCESSORS == //
111impl<V, W: Weight> WeightedList<V,W>
112{
113 /// Get an iterator over copies of the weights of each item in the list.
114 pub fn weights(&self) -> impl Iterator<Item = W>
115 {
116 self.data.iter().map(|item| item.weight)
117 }
118
119 /// Get an iterator over references to the values of each item in the list.
120 pub fn values(&self) -> impl Iterator<Item = &V>
121 {
122 self.data.iter().map(|item| &item.value)
123 }
124
125 /// Get a reference to the `Vec<>` of items in the list.
126 pub fn items(&self) -> &Vec<WeightedItem<V,W>>
127 {
128 &self.data
129 }
130
131 /// Get an iterator over (weight, value) tuples representing each item in the list.
132 ///
133 /// This satisfies the axiom:
134 ///
135 /// ```
136 /// # use weighted_list::*;
137 /// let wl = wlist![(2, "sup"), (3, "nova")];
138 /// let rl = WeightedList::init(wl.raw());
139 ///
140 /// for (left, right) in std::iter::zip(wl.clone(), rl.clone()) {
141 /// assert_eq!(left.weight, right.weight);
142 /// assert_eq!(left.value, *right.value);
143 /// }
144 /// ```
145 pub fn raw(&self) -> impl Iterator<Item = (W,&V)>
146 {
147 self.data.iter().map(|item| (item.weight, &item.value))
148 }
149}
150
151impl<V, W: Weight + nums::PrimInt> WeightedList<V,W>
152{
153 pub fn expanded(&self) -> impl Iterator<Item = &V>
154 {
155 self.data
156 .iter()
157 .flat_map(|item| iter::repeat_n(
158 &item.value,
159 nums::cast::<W, usize>(item.weight).unwrap_or(0)
160 ))
161 }
162}
163
164// == PROPERTIES == //
165impl<V, W: Weight> WeightedList<V,W>
166{
167 /// Sum the weights of all items in the list.
168 ///
169 /// # Notes
170 ///
171 /// - This is not the number of items in the list – use `.total_values()` for that.
172 /// - `len() == 0` does not imply the list is empty – items may have zero or negative weights! To check if the list is empty, use `.is_empty()` instead.
173 pub fn len(&self) -> W
174 {
175 self.data.iter().map(|item| item.weight).sum()
176 }
177
178 pub fn capacity(&self) -> usize
179 {
180 self.data.capacity()
181 }
182
183 /// How many items/values are in the list?
184 ///
185 /// Note that this is not equivalent to `.len()`, which is the total weights of all items in the list.
186 pub fn total_values(&self) -> usize
187 {
188 self.data.len()
189 }
190
191 /// Does the list contain no items?
192 ///
193 /// Note that this returns `false` if the list contains items with weights of `0`.
194 pub fn is_empty(&self) -> bool
195 {
196 self.data.is_empty()
197 }
198
199 /// Do any items have a weight of `0`?
200 pub fn is_zero(&self) -> bool
201 {
202 !self.is_empty()
203 && self.data.iter().all(|item| item.weight == W::zero())
204 }
205
206 /// Do any items have a negative weight?
207 pub fn has_negative_weights(&self) -> bool
208 {
209 !self.is_empty()
210 && self.data.iter().any(|item| item.weight < W::zero())
211 }
212}
213
214// == CONVERSIONS == //
215impl<V, W: Weight> FromIterator<WeightedItem<V,W>> for WeightedList<V,W>
216{
217 fn from_iter<I>(items: I) -> Self
218 where I: IntoIterator<Item = WeightedItem<V,W>>
219 {
220 // TODO benchmark
221 // Self {
222 // data: items.into_iter().collect::<Vec<WeightedItem<V,W>>>()
223 // }
224
225 let mut data = vec![];
226
227 for item in items {
228 data.push(item);
229 }
230
231 Self { data }
232 }
233}
234
235impl<V, W: Weight> From<Vec<WeightedItem<V,W>>> for WeightedList<V,W>
236{
237 fn from(vec: Vec<WeightedItem<V,W>>) -> Self {
238 Self { data: vec }
239 }
240}
241
242impl<V, W: Weight> From<WeightedList<V,W>> for Vec<WeightedItem<V,W>>
243{
244 fn from(list: WeightedList<V,W>) -> Self {
245 list.data
246 }
247}
248
249impl<V, W: Weight> AsRef<Vec<WeightedItem<V,W>>> for WeightedList<V,W>
250{
251 fn as_ref(&self) -> &Vec<WeightedItem<V,W>> {
252 &self.data
253 }
254}
255
256impl<V, W: Weight> Deref for WeightedList<V,W>
257{
258 type Target = [WeightedItem<V,W>];
259
260 fn deref(&self) -> &Self::Target {
261 self.data.deref()
262 }
263}
264
265impl<V, W: Weight> DerefMut for WeightedList<V,W>
266{
267 fn deref_mut(&mut self) -> &mut Self::Target {
268 self.data.deref_mut()
269 }
270}
271
272// == TRAITS == //
273impl<V: fmt::Display, W: Weight + fmt::Display> fmt::Display for WeightedList<V,W>
274{
275 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
276 write!(f,
277 "WeightedList[{}]",
278 self.data.iter().map(|item| item.to_string()).join(", ")
279 )
280 }
281}
282
283impl<V, W: Weight> Extend<WeightedItem<V,W>> for WeightedList<V,W>
284{
285 fn extend<T>(&mut self, iter: T)
286 where T: IntoIterator<Item = WeightedItem<V,W>>
287 {
288 for item in iter {
289 self.push_item(item);
290 }
291 }
292}
293
294// == INTERNAL == //
295impl<V, W: Weight> WeightedList<V,W>
296{
297 /// Convert a `weighted_index` to its unweighted equivalent in the underlying `Vec`. Does not panic on overflow and instead returns the `.len()` of the underlying `Vec`.
298 fn _unweight_index_nopanic_(&self, weighted_index: W) -> usize
299 {
300 let mut t = W::zero();
301 let mut i = 0;
302
303 for item in &self.data {
304 t += item.weight;
305
306 if t > weighted_index {
307 return i;
308 }
309
310 i += 1;
311 }
312
313 i
314 }
315
316 /// Convert a `weighted_index` to its unweighted equivalent in the underlying `Vec`. Panics on overflow.
317 fn _unweight_index_(&self, weighted_index: W) -> usize
318 {
319 let mut t = W::zero();
320
321 for (i, item) in self.data.iter().enumerate() {
322 t += item.weight;
323
324 if t > weighted_index {
325 return i;
326 }
327 }
328
329 panic!(
330 "index out of bounds: the len is {} but the index is {}",
331 self.len(), weighted_index
332 );
333 }
334}
335
336// == INDEXING == //
337impl<V, W: Weight> Index<W> for WeightedList<V,W>
338{
339 type Output = WeightedItem<V,W>;
340
341 fn index(&self, weighted_index: W) -> &Self::Output
342 {
343 let mut t = W::zero();
344
345 for item in &self.data {
346 t += item.weight;
347
348 if t > weighted_index {
349 return item;
350 }
351 };
352
353 panic!("index out of bounds: the len is {} but the index is {weighted_index}", self.len());
354 }
355}
356
357impl<V, W: Weight> IndexMut<W> for WeightedList<V,W>
358{
359 fn index_mut(&mut self, weighted_index: W) -> &mut Self::Output
360 {
361 let idx = self._unweight_index_(weighted_index);
362 &mut self.data[idx]
363 }
364}
365
366// == ITERATION == //
367// TODO: is this still necessary?
368impl<V, W: Weight> WeightedList<V,W>
369{
370 pub fn iter(&self) -> impl Iterator<Item = &WeightedItem<V,W>> {
371 self.data.iter()
372 }
373
374 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut WeightedItem<V,W>> {
375 self.data.iter_mut()
376 }
377}
378
379impl<V, W: Weight> IntoIterator for WeightedList<V,W>
380{
381 type Item = WeightedItem<V,W>;
382 type IntoIter = <Vec<Self::Item> as IntoIterator>::IntoIter;
383
384 /// ```compile_fail
385 /// # use weighted_list::*;
386 /// let list = wlist![]
387 /// for _ in list {}
388 /// list; // compile error
389 /// ```
390 fn into_iter(self) -> Self::IntoIter {
391 self.data.into_iter()
392 }
393}
394
395impl<'l, V, W: Weight> IntoIterator for &'l WeightedList<V,W>
396{
397 type Item = &'l WeightedItem<V,W>;
398 type IntoIter = std::slice::Iter<'l, WeightedItem<V,W>>;
399
400 fn into_iter(self) -> Self::IntoIter {
401 self.data.iter()
402 }
403}
404
405impl<'l, V, W: Weight> IntoIterator for &'l mut WeightedList<V,W>
406{
407 type Item = &'l mut WeightedItem<V,W>;
408 type IntoIter = std::slice::IterMut<'l, WeightedItem<V,W>>;
409
410 fn into_iter(self) -> Self::IntoIter {
411 self.data.iter_mut()
412 }
413}
414
415// == LIST MUTATION == //
416impl<V, W: Weight> WeightedList<V,W>
417{
418 pub fn reserve(&mut self, additional: usize) -> &mut Self
419 {
420 self.data.reserve(additional);
421 self
422 }
423
424 /// Append an item to the end of the list.
425 pub fn push_item(&mut self, item: WeightedItem<V,W>) -> &mut Self
426 {
427 self.data.push(item);
428 self
429 }
430
431 /// Append a new item with `value` and `weight` to the end of the list.
432 pub fn push_new_item(&mut self, weight: W, value: V) -> &mut Self
433 {
434 self.push_item(WeightedItem { weight, value })
435 }
436
437 /// Append a new item with `value` and a weight of `1` to the end of the list.
438 pub fn push_value(&mut self, value: V) -> &mut Self
439 {
440 self.push_item(WeightedItem::unit(value))
441 }
442
443 /// Insert an item into the list at `weighted_index`. If `weighted_index >= len`, the item is appended to the end (the function does *not* panic).
444 pub fn insert_item(&mut self,
445 weighted_index: W,
446 item: WeightedItem<V,W>
447 ) -> &mut Self
448 {
449 self.data.insert(self._unweight_index_nopanic_(weighted_index), item);
450 self
451 }
452
453 /// Insert a new item with `value` and `weight` into the list at `weighted_index`. If `weighted_index >= len`, the item is appended to the end (the function does *not* panic).
454 pub fn insert_new_item(&mut self,
455 weighted_index: W,
456 (weight, value): (W, V)
457 ) -> &mut Self
458 {
459 self.insert_item(weighted_index, WeightedItem::new(weight, value))
460 }
461
462 /// Insert an item with `value` and a weight of `1` into the list at `weighted_index`. If `weighted_index >= len`, the item is appended to the end (the function does *not* panic).
463 pub fn insert_value(&mut self,
464 weighted_index: W,
465 value: V,
466 ) -> &mut Self
467 {
468 self.insert_item(weighted_index, WeightedItem::unit(value))
469 }
470
471 /// Move all items in `other` into `self`, leaving `other` empty.
472 pub fn append(&mut self, other: &mut WeightedList<V,W>) -> &mut Self
473 {
474 self.data.append(&mut other.data);
475 self
476 }
477
478 /// Reverse the order of items in the list.
479 pub fn reverse(&mut self) -> &mut Self
480 {
481 self.data.reverse();
482 self
483 }
484
485 /// Swap the items at weighted indices `left` and `right`.
486 ///
487 /// # Panics
488 ///
489 /// Panics if `left` or `right` are out of bounds.
490 pub fn swap(&mut self, left: W, right: W) -> &mut Self
491 {
492 let l = self._unweight_index_(left);
493 let r = self._unweight_index_(right);
494 self.data.swap(l, r);
495 self
496 }
497
498 /// Removes the last item from the list and returns it, or `None` if the list is empty.
499 pub fn pop(&mut self) -> Option<WeightedItem<V,W>>
500 {
501 self.data.pop()
502 }
503
504 pub fn pop_if(&mut self,
505 predicate: impl FnOnce(&mut WeightedItem<V,W>) -> bool
506 ) -> Option<WeightedItem<V,W>>
507 {
508 self.data.pop_if(predicate)
509 }
510
511 /// Remove the entire item at `weighted_index` and return it.
512 pub fn remove(&mut self, weighted_index: W) -> WeightedItem<V,W>
513 {
514 self.data.remove(self._unweight_index_(weighted_index))
515 }
516
517 // UNTESTED
518 pub fn truncate(&mut self, len: W) -> &mut Self
519 {
520 let mut t = W::zero();
521
522 for each in &mut self.data {
523 t += each.weight;
524
525 if t > len {
526 each.weight = t - len;
527 }
528 }
529
530 self
531 }
532
533 /// Retain only items that fulfil the predicate.
534 pub fn retain<F>(&mut self, predicate: F) -> &mut Self
535 where F: FnMut(&WeightedItem<V,W>) -> bool
536 {
537 self.data.retain(predicate);
538 self
539 }
540
541 /// Retain only items that fulfil the predicate, passing a mutable reference to the predicate.
542 pub fn retain_mut<F>(&mut self, predicate: F) -> &mut Self
543 where F: FnMut(&mut WeightedItem<V,W>) -> bool
544 {
545 self.data.retain_mut(predicate);
546 self
547 }
548
549 /// Clear the list, removing all items.
550 ///
551 /// If you'd like to set the weights of all items to `0`, you can use `.zero_all_weights()`.
552 pub fn clear(&mut self) -> &mut Self
553 {
554 self.data.clear();
555 self
556 }
557}
558
559impl<V: Clone, W: Weight> WeightedList<V,W>
560{
561 /// Return a clone of the list with items sorted in ascending order of weights.
562 ///
563 /// Orderings of items with equivalent weights is (currently) undefined behaviour.
564 pub fn sorted(&self) -> Self
565 where V: Eq, W: Ord
566 {
567 let mut out = self.clone();
568 out.sort();
569 out
570 }
571
572 /// Return a clone of the list with items reversed.
573 pub fn reversed(&self) -> Self
574 {
575 let mut out = self.clone();
576 out.reverse();
577 out
578 }
579}
580
581// == SPECIALISED MUTATION == //
582impl<V, W: Weight> WeightedList<V,W>
583{
584 /// Remove all items with non-positive weight.
585 pub fn prune(&mut self) -> &mut Self
586 {
587 self.data.retain(|item| item.weight > W::zero());
588 self
589 }
590
591 pub fn pruned(&self) -> Self
592 where V: Clone
593 {
594 let mut out = self.clone();
595 out.prune();
596 out
597 }
598
599 /// Set the weight of all items to `0`.
600 pub fn zero_all_weights(&mut self) -> &mut Self
601 {
602 for item in &mut self.data {
603 item.weight = W::zero();
604 }
605
606 self
607 }
608
609 /// Set the weight of all items to `weight`.
610 pub fn set_all_weights(&mut self, weight: W) -> &mut Self
611 {
612 for item in &mut self.data {
613 item.weight = weight;
614 }
615
616 self
617 }
618
619 /// Return a clone of the list with all item weights normalised such that they sum to `1.0`.
620 ///
621 /// # Usage
622 ///
623 /// ```
624 /// # use weighted_list::*;
625 /// let mut wl = wlist![
626 /// (2, "sup"),
627 /// (3, "nova"),
628 /// (5, "shard"),
629 /// ];
630 ///
631 /// assert_eq!(
632 /// wl.normalised().ok(),
633 /// Some(wlist![
634 /// (0.2, "sup"),
635 /// (0.3, "nova"),
636 /// (0.5, "shard"),
637 /// ])
638 /// );
639 /// ```
640 pub fn normalised(&mut self) -> Result<WeightedList<V, f64>, &str>
641 where V: Clone
642 {
643 let total;
644
645 if let Some(t) = nums::cast::<W, f64>(self.len()) {
646 total = t;
647 } else {
648 return Err("Error casting total weights to `f64`");
649 };
650
651 let items =
652 self.data
653 .iter()
654 .map(|item| {
655 let weight = nums::cast::<W, f64>(item.weight)?;
656 Some(WeightedItem {
657 weight: weight / total,
658 value: item.value.clone()
659 })
660 })
661 .collect::<Option<Vec<WeightedItem<V, f64>>>>()
662 ;
663
664 match items {
665 Some(data) => Ok(WeightedList { data }),
666 None => Err("Error casting weights to `f64`")
667 }
668 }
669}
670
671impl<V: PartialEq, W: Weight> WeightedList<V,W>
672{
673 /// Merge a `WeightedItem` into the list. If an item with the same value already exists, add the weight of the new item to the existing item. Otherwise, append the new item to the list.
674 ///
675 /// # Tips
676 ///
677 /// - Use this method when you already have an existing `WeightedItem` instance. If you're going to construct a new `WeightedItem`, `.merge_new_item()` will be more convenient.
678 ///
679 /// # Usage
680 ///
681 /// ```
682 /// # use weighted_list::*;
683 /// let mut wl = wlist![(1, "sup")];
684 ///
685 /// let item = WeightedItem::new(2, "sup");
686 /// wl.merge_item(item);
687 /// assert!(wl == wlist![(3, "sup")]);
688 /// // "sup" merged with existing instance
689 ///
690 /// let item = WeightedItem::unit("elysion");
691 /// wl.merge_item(item);
692 /// assert!(wl == wlist![(3, "sup"), (1, "elysion")]);
693 /// // "elysion" appended to end
694 /// ```
695 pub fn merge_item(&mut self, item: WeightedItem<V,W>) -> &mut Self
696 {
697 if let Some(existing) = self.data.iter_mut().find(|each| each.value == item.value) {
698 existing.weight += item.weight;
699 }
700 else {
701 self.data.push(item);
702 }
703
704 self
705 }
706
707 /// Merge a new item with `value` and `weight` into the list.
708 ///
709 /// See `.merge_item()` for details.
710 pub fn merge_new_item(&mut self, weight: W, value: V) -> &mut Self
711 {
712 self.merge_item(WeightedItem { weight, value })
713 }
714
715 /// Merge a new item with `value` and a weight of `1` into the list.
716 ///
717 /// See `.merge_item()` for details.
718 pub fn merge_value(&mut self, value: V) -> &mut Self
719 {
720 self.merge_item(WeightedItem::unit(value))
721 }
722
723 /// Merge the items of `other` into `self`, leaving `other` empty.
724 ///
725 /// See `.merge_item()` for details.
726 pub fn merge_with(&mut self, other: WeightedList<V,W>) -> &mut Self
727 {
728 for item in other {
729 self.merge_item(item);
730 }
731
732 self
733 }
734
735 /// Merge any duplicate items in the list.
736 ///
737 /// # Usage
738 ///
739 /// ```
740 /// # use weighted_list::*;
741 /// let mut wl = wlist![
742 /// (1, "sup"),
743 /// (2, ""),
744 /// (20, "sup")
745 /// ];
746 ///
747 /// assert_eq!(
748 /// *wl.merge_duplicates(),
749 /// wlist![
750 /// (21, "sup"),
751 /// (2, ""),
752 /// ]
753 /// );
754 /// ```
755 pub fn merge_duplicates(&mut self) -> &mut Self
756 {
757 let orig = std::mem::replace(self, WeightedList::new());
758 self.merge_with(orig);
759 self
760 }
761}
762
763impl<V: Clone, W: Weight> WeightedList<V,W>
764{
765 /// Decrement the weight of the item at `weighted_index` by `1`. If its weight becomes non-positive as a result, remove the entire item. Returns a clone of the item with its updated weight.
766 pub fn take_one(&mut self, weighted_index: W) -> WeightedItem<V,W>
767 {
768 self.take_by(weighted_index, W::one())
769 }
770
771 /// Decrement the weight of the item at `weighted_index` by `decrement`. If its weight becomes non-positive as a result, remove the entire item. Returns a clone of the item with its updated weight.
772 pub fn take_by(&mut self, weighted_index: W, decrement: W) -> WeightedItem<V,W>
773 {
774 let idx = self._unweight_index_(weighted_index);
775 let target = &mut self.data[idx];
776
777 target.weight -= decrement;
778
779 if target.weight <= W::zero() {
780 self.data.remove(idx)
781 }
782 else {
783 target.clone()
784 }
785 }
786
787 /// Remove the entire item at `weighted_index`.
788 pub fn take_entire(&mut self, weighted_index: W) -> WeightedItem<V,W>
789 {
790 self.remove(weighted_index)
791 }
792}
793
794// == RANDOMISATION == //
795impl<V, W: Weight> WeightedList<V,W>
796{
797 fn _get_random_weighted_index_up_to_<RNG>(&self, rng: &mut RNG, upper: W) -> Option<W>
798 where RNG: Rng + ?Sized
799 {
800 let len: f64 = nums::cast::<W, f64>(upper)?;
801 let scalar: f64 = rng.random();
802
803 let idx = (len * scalar).floor();
804 let out = nums::cast::<f64, W>(idx)?;
805
806 Some(out)
807 }
808
809 fn _get_random_weighted_index_<RNG>(&self, rng: &mut RNG) -> Option<W>
810 where RNG: Rng + ?Sized
811 {
812 self._get_random_weighted_index_up_to_(rng, self.len())
813 }
814
815 /// Select a random item from the list and return its value, using weighted randomisation.
816 ///
817 /// # Usage
818 ///
819 /// ```
820 /// # use weighted_list::*;
821 /// let wl = wlist![
822 /// (2, String::from("sup")),
823 /// (3, String::from("nova")),
824 /// (5, String::from("shard")),
825 /// ];
826 ///
827 /// wl.select_random_value(&mut rand::rng());
828 /// // could give:
829 /// // - Some("sup" ) with 20% probability
830 /// // - Some("nova" ) with 30% probability
831 /// // - Some("shard") with 50% probability
832 /// ```
833 pub fn select_random_value<RNG>(&self, rng: &mut RNG) -> Option<&V>
834 where RNG: Rng + ?Sized
835 {
836 let out = &self.select_random_item(rng)?.value;
837 Some(out)
838 }
839
840 /// Select a random item from the list, using weighted randomisation.
841 ///
842 /// # Usage
843 ///
844 /// ```
845 /// # use weighted_list::*;
846 /// let wl = wlist![
847 /// (2, String::from("sup")),
848 /// (3, String::from("nova")),
849 /// (5, String::from("shard")),
850 /// ];
851 ///
852 /// wl.select_random_item(&mut rand::rng());
853 /// // could give:
854 /// // - Some(WeightedItem { 2, "sup" }) with 20% probability
855 /// // - Some(WeightedItem { 3, "nova" }) with 30% probability
856 /// // - Some(WeightedItem { 5, "shard" }) with 50% probability
857 /// ```
858 pub fn select_random_item<RNG>(&self, rng: &mut RNG) -> Option<&WeightedItem<V,W>>
859 where RNG: Rng + ?Sized
860 {
861 if self.data.is_empty() { return None }
862
863 let idx = self._get_random_weighted_index_(rng)?;
864 let out = &self[idx];
865
866 Some(out)
867 }
868}
869
870impl<V: Clone, W: Weight> WeightedList<V,W>
871{
872 /// Select a random item from the list using weighted randomisation, and decrement its weight by `1`.
873 ///
874 /// # Usage
875 ///
876 /// ```
877 /// # use weighted_list::*;
878 /// let mut wl = wlist![
879 /// (2, String::from("sup")),
880 /// (3, String::from("nova")),
881 /// (5, String::from("shard")),
882 /// ];
883 ///
884 /// wl.take_one_random(&mut rand::rng());
885 /// // could give:
886 /// // - Some(WeightedItem { 1, "sup" }) with 20% probability
887 /// // - Some(WeightedItem { 2, "nova" }) with 30% probability
888 /// // - Some(WeightedItem { 4, "shard" }) with 50% probability
889 /// ```
890 pub fn take_one_random<RNG>(&mut self, rng: &mut RNG) -> Option<WeightedItem<V,W>>
891 where RNG: Rng + ?Sized
892 {
893 self.take_by_random(rng, W::one())
894 }
895
896 /// Select a random item from the list using weighted randomisation, and decrement its weight by `decrement`.
897 ///
898 /// # Usage
899 ///
900 /// ```
901 /// # use weighted_list::*;
902 /// let mut wl = wlist![
903 /// (2, String::from("sup")),
904 /// (3, String::from("nova")),
905 /// (5, String::from("shard")),
906 /// ];
907 ///
908 /// wl.take_by_random(&mut rand::rng(), 2);
909 /// // could give:
910 /// // - Some(WeightedItem { 0, "sup" }) with 20% probability
911 /// // - Some(WeightedItem { 1, "nova" }) with 30% probability
912 /// // - Some(WeightedItem { 3, "shard" }) with 50% probability
913 /// ```
914 pub fn take_by_random<RNG>(&mut self, rng: &mut RNG, decrement: W) -> Option<WeightedItem<V,W>>
915 where RNG: Rng + ?Sized
916 {
917 if self.data.is_empty() { return None }
918
919 let idx = self._get_random_weighted_index_(rng)?;
920 let out = self.take_by(idx, decrement);
921
922 Some(out)
923 }
924
925 /// Select and remove a random item from the list, using weighted randomisation.
926 ///
927 /// # Usage
928 ///
929 /// ```
930 /// # use weighted_list::*;
931 /// let mut wl = wlist![
932 /// (2, String::from("sup")),
933 /// (3, String::from("nova")),
934 /// (5, String::from("shard")),
935 /// ];
936 ///
937 /// wl.take_entire_random(&mut rand::rng());
938 /// // could give:
939 /// // - Some(WeightedItem { 2, "sup" }) with 20% probability
940 /// // - Some(WeightedItem { 3, "nova" }) with 30% probability
941 /// // - Some(WeightedItem { 5, "shard" }) with 50% probability
942 ///
943 /// assert!( wl.total_values() == 2 );
944 /// ```
945 pub fn take_entire_random<RNG>(&mut self, rng: &mut RNG) -> Option<WeightedItem<V,W>>
946 where RNG: Rng + ?Sized
947 {
948 if self.data.is_empty() { return None }
949
950 let idx = self._get_random_weighted_index_(rng)?;
951 let out = self.take_entire(idx);
952
953 Some(out)
954 }
955}
956
957#[bon]
958impl<V: Clone + Eq, W: Weight> WeightedList<V,W>
959{
960 /// Select `count` values using weighted randomisation.
961 ///
962 /// Call this method using `bon` builder syntax (see § Usage below).
963 ///
964 /// # Options
965 ///
966 /// ```ignore
967 /// rng: RNG,
968 /// count: usize,
969 /// replace: Option<bool> = true,
970 /// decrement: Option<W> = 1,
971 /// unique: Option<bool> = false,
972 /// ```
973 ///
974 /// - `count`: How many values to select.
975 /// - `replace`: If `true`, items do not have their weight decremented after selection, and infinite values can be selected. If `false`, items have their weight decremented after selection.
976 /// - This means at most `self.len()` values will be returned.
977 /// - `decrement`: How much to decrement weights by if `replace` is `false`.
978 /// - `unique`: If `true`, only distinct values will be returned.
979 /// - `replace` becomes irrelevant in this case.
980 /// - This uses `Eq` equality comparison.
981 /// - This means at most `self.total_values()` values will be returned.
982 ///
983 /// # Usage
984 ///
985 /// This method uses the bon builder syntax:
986 ///
987 /// ```
988 /// # use weighted_list::*;
989 /// let mut pool = wlist![
990 /// (2, String::from("sup")),
991 /// (3, String::from("nova")),
992 /// (5, String::from("shard")),
993 /// ];
994 ///
995 /// let mut rng = rand::rng();
996 ///
997 /// // with replacement
998 /// let selected =
999 /// pool.select_random_values()
1000 /// .rng(&mut rng)
1001 /// .count(3)
1002 /// .call();
1003 ///
1004 /// assert!(selected.len() == 3);
1005 ///
1006 /// // without replacement
1007 /// let selected =
1008 /// pool.select_random_values()
1009 /// .rng(&mut rng)
1010 /// .count(10)
1011 /// .replace(false)
1012 /// .decrement(2)
1013 /// .call();
1014 ///
1015 /// assert!(selected.len() == 6);
1016 ///
1017 /// // unique only
1018 /// let selected =
1019 /// pool.select_random_values()
1020 /// .rng(&mut rng)
1021 /// .count(100)
1022 /// .unique(true)
1023 /// .call();
1024 ///
1025 /// assert!(selected.len() == 3);
1026 /// ```
1027 ///
1028 /// # Notes
1029 ///
1030 /// - It is not guaranteed that the results will have exactly `count` values.
1031 /// - If `count` exceeds the maximum possible number of values that can be returned, excess iterations will be skipped.
1032 /// - If selection for an iteration fails, that value is excluded from the output list.
1033 /// - This method reserves a `Vec<>` with capacity `count` initially, so be careful of passing in extremely large `count`s.
1034 #[builder]
1035 pub fn select_random_values<RNG>(&self,
1036 rng: &mut RNG,
1037 count: usize,
1038 replace: Option<bool>,
1039 decrement: Option<W>,
1040 unique: Option<bool>,
1041 ) -> Vec<V>
1042 where RNG: Rng + ?Sized
1043 {
1044 let replace = replace.unwrap_or(true);
1045 let decrement = decrement.unwrap_or(W::one());
1046 let unique = unique.unwrap_or(false);
1047
1048 let mut pool = self.clone();
1049 let mut n = 0;
1050 let mut out = Vec::with_capacity(count);
1051
1052 loop
1053 {
1054 n += 1;
1055 if n > count || self.data.is_empty() { break }
1056
1057 if let Some(item) = {
1058 if unique { pool.take_entire_random(rng) }
1059 else if replace { pool.take_by_random(rng, W::zero()) }
1060 else { pool.take_by_random(rng, decrement) }
1061 } {
1062 out.push(item.value.clone());
1063 }
1064 }
1065
1066 out
1067 }
1068
1069 /// Take `count` values using weighted randomisation.
1070 #[builder]
1071 pub fn take_random_values<RNG>(&mut self,
1072 rng: &mut RNG,
1073 count: usize,
1074 take_entire: Option<bool>,
1075 decrement: Option<W>,
1076 ) -> Vec<V>
1077 where RNG: Rng + ?Sized
1078 {
1079 let take_entire = take_entire.unwrap_or(true);
1080 let decrement = decrement.unwrap_or(W::one());
1081
1082 let mut n = 0;
1083 let mut out = Vec::with_capacity(count as usize);
1084
1085 loop
1086 {
1087 n += 1;
1088 if n > count || self.data.is_empty() { break }
1089
1090 if let Some(item) = {
1091 if take_entire { self.take_entire_random(rng) }
1092 else { self.take_by_random(rng, decrement) }
1093 } {
1094 out.push(item.value.clone());
1095 }
1096 }
1097
1098 out
1099 }
1100}
1101
1102#[bon]
1103impl<V: Clone + Eq + std::hash::Hash, W: Weight> WeightedList<V,W>
1104{
1105 /// Variant of `._unweighted_index_()` for mutating random selection with unique outputs.
1106 fn _unweight_index_skipping_(&self,
1107 weighted_index: W,
1108 seen: &std::collections::HashSet<V>,
1109 ) -> Option<usize>
1110 {
1111 let mut t = W::zero();
1112
1113 for (i, item) in self.data.iter().enumerate() {
1114 if seen.contains(&item.value) {
1115 continue
1116 }
1117
1118 t += item.weight;
1119
1120 if t > weighted_index {
1121 return Some(i);
1122 }
1123 }
1124
1125 None
1126 }
1127
1128 /// Take `count` values using weighted randomisation.
1129 #[builder]
1130 pub fn take_random_values_unique<RNG>(&mut self,
1131 rng: &mut RNG,
1132 count: usize,
1133 decrement: Option<W>,
1134 ) -> Vec<V>
1135 where RNG: Rng + ?Sized,
1136 {
1137 let decrement = decrement.unwrap_or(W::one());
1138
1139 let mut n = 0;
1140 let mut l = self.len();
1141
1142 let mut seen = std::collections::HashSet::<V>::new();
1143
1144 let mut out = Vec::with_capacity(
1145 if count > 16 {
1146 count.min(self.total_values())
1147 } else {
1148 count
1149 }
1150 );
1151
1152 loop
1153 {
1154 n += 1;
1155 if n > count || self.data.is_empty() { break }
1156
1157 if let Some(value) = (|| {
1158 let weighted_index = self._get_random_weighted_index_up_to_(rng, l)?;
1159 let idx = self._unweight_index_skipping_(weighted_index, &seen)?;
1160
1161 let target = &mut self.data[idx];
1162 let value = target.value.clone();
1163
1164 target.weight -= decrement;
1165
1166 if target.weight <= W::zero() {
1167 self.data.remove(idx);
1168 }
1169
1170 Some(value)
1171 })()
1172 {
1173 seen.insert(value.clone());
1174 out.push(value.clone());
1175
1176 l = self.data.iter()
1177 .filter_map(
1178 |item| {
1179 if !seen.contains(&item.value) { Some(item.weight) }
1180 else { None }
1181 }
1182 )
1183 .sum::<W>();
1184 }
1185 }
1186
1187 out
1188 }
1189}
1190
1191impl<V: Clone, W: Weight + Clone> WeightedList<V,W>
1192{
1193 /// Shuffle the order of items in the list in-place.
1194 pub fn shuffle_items<RNG>(&mut self, rng: &mut RNG) -> &mut Self
1195 where RNG: Rng + ?Sized
1196 {
1197 self.data.shuffle(rng);
1198 self
1199 }
1200
1201 /// Return a clone with the order of items shuffled.
1202 ///
1203 /// Out-of-place version of `.shuffle_items()`.
1204 pub fn shuffled_items<RNG>(&self, rng: &mut RNG) -> Self
1205 where RNG: Rng + ?Sized
1206 {
1207 let mut out = self.clone();
1208 out.shuffle_items(rng);
1209
1210 out
1211 }
1212
1213 /// Shuffle the pairings of (weight, value) for items in the list, in-place.
1214 ///
1215 /// Values remain in the same order, while weights are re-assigned.
1216 ///
1217 /// # Usage
1218 ///
1219 /// ```
1220 /// # use weighted_list::*;
1221 /// let mut wl = wlist![
1222 /// (2, "sup"),
1223 /// (3, "nova"),
1224 /// (5, "shard"),
1225 /// ];
1226 ///
1227 /// wl.shuffle_weights(&mut rand::rng());
1228 ///
1229 /// println!("{wl}");
1230 /// // could give:
1231 /// // WeightedList[{3, sup}, {5, nova}, {2, shard}]
1232 ///
1233 /// println!("{wl}");
1234 /// // could now give:
1235 /// // WeightedList[{2, sup}, {5, nova}, {3, shard}]
1236 /// ```
1237 pub fn shuffle_weights<RNG>(&mut self, rng: &mut RNG) -> &mut Self
1238 where RNG: Rng + ?Sized
1239 {
1240 let mut weights: Vec<W> = self.weights().collect();
1241 weights.shuffle(rng);
1242
1243 for item in &mut self.data {
1244 /* guaranteed to be Some */
1245 item.weight = weights.pop().unwrap();
1246 }
1247
1248 self
1249 }
1250
1251 /// Return a clone of the list with (weight, value) pairings shuffled.
1252 ///
1253 /// Out-of-place version of `.shuffle_weights()`.
1254 pub fn shuffled_weights<RNG>(&self, rng: &mut RNG) -> Self
1255 where RNG: Rng + ?Sized
1256 {
1257 let mut out = self.clone();
1258 out.shuffle_weights(rng);
1259
1260 out
1261 }
1262}
1263
1264
1265// == INTERNAL TESTS == //
1266#[cfg(test)]
1267mod tests
1268{
1269 use super::*;
1270
1271 fn el() -> WeightedList<String, i32>
1272 {
1273 wlist![]
1274 }
1275
1276 fn wl() -> WeightedList<String, i32>
1277 {
1278 wlist![
1279 (2, "sup".to_string()),
1280 (3, "nova".to_string()),
1281 (5, "shard".to_string()),
1282 ]
1283 }
1284
1285 #[test] fn _unweight_index_()
1286 {
1287 let list = wl();
1288 assert_eq!( list._unweight_index_(0), 0 );
1289 assert_eq!( list._unweight_index_(1), 0 );
1290 assert_eq!( list._unweight_index_(2), 1 );
1291 assert_eq!( list._unweight_index_(3), 1 );
1292 assert_eq!( list._unweight_index_(4), 1 );
1293 assert_eq!( list._unweight_index_(5), 2 );
1294 assert_eq!( list._unweight_index_(6), 2 );
1295 assert_eq!( list._unweight_index_(7), 2 );
1296 assert_eq!( list._unweight_index_(8), 2 );
1297 assert_eq!( list._unweight_index_(9), 2 );
1298 }
1299
1300 #[test] #[should_panic] fn _unweight_index_empty_()
1301 {
1302 el()._unweight_index_(0);
1303 }
1304
1305 #[test] #[should_panic] fn _unweight_index_out_of_bounds_()
1306 {
1307 wl()._unweight_index_(10);
1308 }
1309
1310 #[test] fn _unweight_index_nopanic_()
1311 {
1312 let list = wl();
1313 assert_eq!( list._unweight_index_nopanic_(10), 3 );
1314 assert_eq!( list._unweight_index_nopanic_(11), 3 );
1315 assert_eq!( list._unweight_index_nopanic_(12), 3 );
1316 }
1317
1318 #[test] fn _unweight_index_skipping_()
1319 {
1320 let list = wl();
1321
1322 let seen = std::collections::HashSet::from(["sup".to_string()]);
1323 assert_eq!( list._unweight_index_skipping_(0, &seen), Some(1) );
1324 assert_eq!( list._unweight_index_skipping_(1, &seen), Some(1) );
1325 assert_eq!( list._unweight_index_skipping_(2, &seen), Some(1) );
1326 assert_eq!( list._unweight_index_skipping_(3, &seen), Some(2) );
1327 assert_eq!( list._unweight_index_skipping_(4, &seen), Some(2) );
1328 assert_eq!( list._unweight_index_skipping_(5, &seen), Some(2) );
1329 assert_eq!( list._unweight_index_skipping_(6, &seen), Some(2) );
1330 assert_eq!( list._unweight_index_skipping_(7, &seen), Some(2) );
1331
1332 let seen = std::collections::HashSet::from(["nova".to_string()]);
1333 assert_eq!( list._unweight_index_skipping_(0, &seen), Some(0) );
1334 assert_eq!( list._unweight_index_skipping_(1, &seen), Some(0) );
1335 assert_eq!( list._unweight_index_skipping_(2, &seen), Some(2) );
1336 assert_eq!( list._unweight_index_skipping_(3, &seen), Some(2) );
1337 assert_eq!( list._unweight_index_skipping_(4, &seen), Some(2) );
1338 assert_eq!( list._unweight_index_skipping_(5, &seen), Some(2) );
1339 assert_eq!( list._unweight_index_skipping_(6, &seen), Some(2) );
1340 }
1341}