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