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// == INDEXING == //
295impl<V, W: Weight> Index<W> for WeightedList<V,W>
296{
297 type Output = WeightedItem<V,W>;
298
299 fn index(&self, weighted_index: W) -> &Self::Output
300 {
301 let mut t = W::zero();
302
303 for item in &self.data {
304 t += item.weight;
305
306 if t > weighted_index {
307 return item;
308 }
309 };
310
311 panic!("index out of bounds: the len is {} but the index is {weighted_index}", self.len());
312 }
313}
314
315impl<V, W: Weight> IndexMut<W> for WeightedList<V,W>
316{
317 fn index_mut(&mut self, weighted_index: W) -> &mut Self::Output
318 {
319 let idx = self._unweight_index_(weighted_index);
320 &mut self.data[idx]
321 }
322}
323
324// == ITERATION == //
325impl<V, W: Weight> IntoIterator for WeightedList<V,W>
326{
327 type Item = WeightedItem<V,W>;
328 type IntoIter = <Vec<Self::Item> as IntoIterator>::IntoIter;
329
330 /// ```compile_fail
331 /// # use weighted_list::*;
332 /// let list = wlist![]
333 /// for _ in list {}
334 /// list; // compile error
335 /// ```
336 fn into_iter(self) -> Self::IntoIter {
337 self.data.into_iter()
338 }
339}
340
341impl<'l, V, W: Weight> IntoIterator for &'l WeightedList<V,W>
342{
343 type Item = &'l WeightedItem<V,W>;
344 type IntoIter = std::slice::Iter<'l, WeightedItem<V,W>>;
345
346 fn into_iter(self) -> Self::IntoIter {
347 self.data.iter()
348 }
349}
350
351impl<'l, V, W: Weight> IntoIterator for &'l mut WeightedList<V,W>
352{
353 type Item = &'l mut WeightedItem<V,W>;
354 type IntoIter = std::slice::IterMut<'l, WeightedItem<V,W>>;
355
356 fn into_iter(self) -> Self::IntoIter {
357 self.data.iter_mut()
358 }
359}
360
361// == LIST MUTATION == //
362impl<V, W: Weight> WeightedList<V,W>
363{
364 pub fn reserve(&mut self, additional: usize) -> &mut Self
365 {
366 self.data.reserve(additional);
367 self
368 }
369
370 /// Append an item to the end of the list.
371 ///
372 /// # Usage
373 ///
374 /// ```
375 /// # use weighted_list::*;
376 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
377 ///
378 /// assert_eq!(
379 /// *wl.push_item(wl[0].clone()),
380 /// wlist![(2, "sup"), (3, "nova"), (5, "shard"), (2, "sup")],
381 /// )
382 /// ```
383 pub fn push_item(&mut self, item: WeightedItem<V,W>) -> &mut Self
384 {
385 self.data.push(item);
386 self
387 }
388
389 /// Append a new item with `value` and `weight` to the end of the list.
390 ///
391 /// # Usage
392 ///
393 /// ```
394 /// # use weighted_list::*;
395 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
396 ///
397 /// assert_eq!(
398 /// *wl.push_new_item(7, "cortex"),
399 /// wlist![(2, "sup"), (3, "nova"), (5, "shard"), (7, "cortex")],
400 /// )
401 /// ```
402 pub fn push_new_item(&mut self, weight: W, value: V) -> &mut Self
403 {
404 self.push_item(WeightedItem { weight, value })
405 }
406
407 /// Append a new item with `value` and a weight of `1` to the end of the list.
408 ///
409 /// # Usage
410 ///
411 /// ```
412 /// # use weighted_list::*;
413 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
414 ///
415 /// assert_eq!(
416 /// *wl.push_value("elysion"),
417 /// wlist![(2, "sup"), (3, "nova"), (5, "shard"), (1, "elysion")],
418 /// )
419 /// ```
420 pub fn push_value(&mut self, value: V) -> &mut Self
421 {
422 self.push_item(WeightedItem::unit(value))
423 }
424
425 /// 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).
426 pub fn insert_item(&mut self,
427 weighted_index: W,
428 item: WeightedItem<V,W>
429 ) -> &mut Self
430 {
431 self.data.insert(self._unweight_index_nopanic_(weighted_index), item);
432 self
433 }
434
435 /// 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).
436 pub fn insert_new_item(&mut self,
437 weighted_index: W,
438 (weight, value): (W, V)
439 ) -> &mut Self
440 {
441 self.insert_item(weighted_index, WeightedItem::new(weight, value))
442 }
443
444 /// 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).
445 pub fn insert_value(&mut self,
446 weighted_index: W,
447 value: V,
448 ) -> &mut Self
449 {
450 self.insert_item(weighted_index, WeightedItem::unit(value))
451 }
452
453 /// Move all items in `other` into `self`, leaving `other` empty.
454 pub fn append(&mut self, other: &mut WeightedList<V,W>) -> &mut Self
455 {
456 self.data.append(&mut other.data);
457 self
458 }
459
460 /// Reverse the order of items in the list.
461 pub fn reverse(&mut self) -> &mut Self
462 {
463 self.data.reverse();
464 self
465 }
466
467 /// Swap the items at weighted indices `left` and `right`.
468 ///
469 /// # Panics
470 ///
471 /// Panics if `left` or `right` are out of bounds.
472 pub fn swap(&mut self, left: W, right: W) -> &mut Self
473 {
474 let l = self._unweight_index_(left);
475 let r = self._unweight_index_(right);
476 self.data.swap(l, r);
477 self
478 }
479
480 /// Removes the last item from the list and returns it, or `None` if the list is empty.
481 pub fn pop(&mut self) -> Option<WeightedItem<V,W>>
482 {
483 self.data.pop()
484 }
485
486 pub fn pop_if(&mut self,
487 predicate: impl FnOnce(&mut WeightedItem<V,W>) -> bool
488 ) -> Option<WeightedItem<V,W>>
489 {
490 self.data.pop_if(predicate)
491 }
492
493 /// Remove the entire item at `weighted_index` and return it.
494 ///
495 /// # Panics
496 ///
497 /// Panics if `weighted_index` is out of bounds.
498 pub fn remove_at(&mut self, weighted_index: W) -> WeightedItem<V,W>
499 {
500 self.data.remove(self._unweight_index_(weighted_index))
501 }
502
503 // UNTESTED
504 pub fn truncate(&mut self, len: W) -> &mut Self
505 {
506 let mut t = W::zero();
507
508 for each in &mut self.data {
509 t += each.weight;
510
511 if t > len {
512 each.weight = t - len;
513 }
514 }
515
516 self
517 }
518
519 /// Retain only items that fulfil the predicate.
520 pub fn retain<F>(&mut self, predicate: F) -> &mut Self
521 where F: FnMut(&WeightedItem<V,W>) -> bool
522 {
523 self.data.retain(predicate);
524 self
525 }
526
527 /// Retain only items that fulfil the predicate, passing a mutable reference to the predicate.
528 pub fn retain_mut<F>(&mut self, predicate: F) -> &mut Self
529 where F: FnMut(&mut WeightedItem<V,W>) -> bool
530 {
531 self.data.retain_mut(predicate);
532 self
533 }
534
535 /// Clear the list, removing all items.
536 ///
537 /// If you'd like to set the weights of all items to `0`, you can use `.zero_all_weights()`.
538 pub fn clear(&mut self) -> &mut Self
539 {
540 self.data.clear();
541 self
542 }
543}
544
545impl<V: Clone, W: Weight> WeightedList<V,W>
546{
547 /// Return a clone of the list with items sorted in ascending order of weights.
548 ///
549 /// Orderings of items with equivalent weights is (currently) undefined behaviour.
550 pub fn sorted(&self) -> Self
551 where V: Eq, W: Ord
552 {
553 let mut out = self.clone();
554 out.sort();
555 out
556 }
557
558 /// Return a clone of the list with items reversed.
559 pub fn reversed(&self) -> Self
560 {
561 let mut out = self.clone();
562 out.reverse();
563 out
564 }
565}
566
567// == SPECIALISED MUTATION == //
568impl<V, W: Weight> WeightedList<V,W>
569{
570 /// Remove all items with non-positive weight.
571 pub fn prune(&mut self) -> &mut Self
572 {
573 self.data.retain(|item| item.weight > W::zero());
574 self
575 }
576
577 pub fn pruned(&self) -> Self
578 where V: Clone
579 {
580 let mut out = self.clone();
581 out.prune();
582 out
583 }
584
585 /// Find the first occurrence (from the left) of an item with `value`, and remove the entire item.
586 ///
587 /// # Usage
588 ///
589 /// ```
590 /// # use weighted_list::*;
591 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
592 ///
593 /// assert_eq!(
594 /// *wl.remove_value_first(&"nova"),
595 /// wlist![(2, "sup"), (5, "shard")],
596 /// )
597 /// ```
598 pub fn remove_value_first(&mut self, value: &V) -> &mut Self
599 where V: PartialEq
600 {
601 self.remove_first_where(|item| item.value == *value)
602 }
603
604 /// Find the last occurrence (from the right) of an item with `value`, and remove the entire item.
605 ///
606 /// # Usage
607 ///
608 /// ```
609 /// # use weighted_list::*;
610 /// let mut wl = wlist![(0, "qi"), (1, "qi"), (2, "sup")];
611 ///
612 /// assert_eq!(
613 /// *wl.remove_value_last(&"qi"),
614 /// wlist![(0, "qi"), (2, "sup")],
615 /// )
616 /// ```
617 pub fn remove_value_last(&mut self, value: &V) -> &mut Self
618 where V: PartialEq
619 {
620 self.remove_last_where(|item| item.value == *value)
621 }
622
623 /// Find the first occurrence (from the left) of an item that matches `predicate`, and remove the entire item.
624 ///
625 /// # Usage
626 ///
627 /// ```
628 /// # use weighted_list::*;
629 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
630 ///
631 /// assert_eq!(
632 /// *wl.remove_first_where(|item| item.weight > 2),
633 /// wlist![(2, "sup"), (5, "shard")],
634 /// )
635 /// ```
636 pub fn remove_first_where<F>(&mut self, predicate: F) -> &mut Self
637 where F: FnMut(&WeightedItem<V,W>) -> bool
638 {
639 if let Some(idx) = self.iter().position(predicate) {
640 self.data.remove(idx);
641 }
642
643 self
644 }
645
646 /// Find the last occurrence (from the right) of an item that matches `predicate`, and remove the entire item.
647 ///
648 /// # Usage
649 ///
650 /// ```
651 /// # use weighted_list::*;
652 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
653 ///
654 /// assert_eq!(
655 /// *wl.remove_last_where(|item| item.weight > 2),
656 /// wlist![(2, "sup"), (3, "nova")],
657 /// );
658 /// ```
659 pub fn remove_last_where<F>(&mut self, predicate: F) -> &mut Self
660 where F: FnMut(&WeightedItem<V,W>) -> bool
661 {
662 if let Some(idx) = self.iter().rposition(predicate) {
663 self.data.remove(idx);
664 }
665
666 self
667 }
668
669 /// Set the weight of all items to `0`.
670 ///
671 /// # Usage
672 ///
673 /// ```
674 /// # use weighted_list::*;
675 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
676 ///
677 /// assert_eq!(
678 /// *wl.zero_all_weights(),
679 /// wlist![(0, "sup"), (0, "nova"), (0, "shard")],
680 /// )
681 /// ```
682 pub fn zero_all_weights(&mut self) -> &mut Self
683 {
684 for item in &mut self.data {
685 item.weight = W::zero();
686 }
687
688 self
689 }
690
691 /// Set the weight of all items to `weight`.
692 ///
693 /// # Usage
694 ///
695 /// ```
696 /// # use weighted_list::*;
697 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
698 ///
699 /// assert_eq!(
700 /// *wl.set_all_weights(1),
701 /// wlist![(1, "sup"), (1, "nova"), (1, "shard")],
702 /// )
703 /// ```
704 pub fn set_all_weights(&mut self, weight: W) -> &mut Self
705 {
706 for item in &mut self.data {
707 item.weight = weight;
708 }
709
710 self
711 }
712
713 /// Return a clone of the list with all item weights normalised such that they sum to `1.0`.
714 ///
715 /// # Usage
716 ///
717 /// ```
718 /// # use weighted_list::*;
719 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
720 ///
721 /// assert_eq!(
722 /// wl.normalised().ok(),
723 /// Some(wlist![(0.2, "sup"), (0.3, "nova"), (0.5, "shard")])
724 /// );
725 /// ```
726 pub fn normalised(&mut self) -> Result<WeightedList<V, f64>, &str>
727 where V: Clone
728 {
729 let total;
730
731 if let Some(t) = nums::cast::<W, f64>(self.len()) {
732 total = t;
733 } else {
734 return Err("Error casting total weights to `f64`");
735 };
736
737 let items =
738 self.data
739 .iter()
740 .map(|item| {
741 let weight = nums::cast::<W, f64>(item.weight)?;
742 Some(WeightedItem {
743 weight: weight / total,
744 value: item.value.clone()
745 })
746 })
747 .collect::<Option<Vec<WeightedItem<V, f64>>>>()
748 ;
749
750 match items {
751 Some(data) => Ok(WeightedList { data }),
752 None => Err("Error casting weights to `f64`")
753 }
754 }
755}
756
757impl<V: PartialEq, W: Weight> WeightedList<V,W>
758{
759 /// 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.
760 ///
761 /// # Tips
762 ///
763 /// - 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.
764 ///
765 /// # Usage
766 ///
767 /// ```
768 /// # use weighted_list::*;
769 /// let mut wl = wlist![(1, "sup")];
770 ///
771 /// let item = WeightedItem::new(2, "sup");
772 /// wl.merge_item(item);
773 /// assert!(wl == wlist![(3, "sup")]);
774 /// // "sup" merged with existing instance
775 ///
776 /// let item = WeightedItem::unit("elysion");
777 /// wl.merge_item(item);
778 /// assert!(wl == wlist![(3, "sup"), (1, "elysion")]);
779 /// // "elysion" appended to end
780 /// ```
781 pub fn merge_item(&mut self, item: WeightedItem<V,W>) -> &mut Self
782 {
783 if let Some(existing) = self.data.iter_mut().find(|each| each.value == item.value) {
784 existing.weight += item.weight;
785 }
786 else {
787 self.data.push(item);
788 }
789
790 self
791 }
792
793 /// Merge a new item with `value` and `weight` into the list.
794 ///
795 /// See `.merge_item()` for details.
796 pub fn merge_new_item(&mut self, weight: W, value: V) -> &mut Self
797 {
798 self.merge_item(WeightedItem { weight, value })
799 }
800
801 /// Merge a new item with `value` and a weight of `1` into the list.
802 ///
803 /// See `.merge_item()` for details.
804 pub fn merge_value(&mut self, value: V) -> &mut Self
805 {
806 self.merge_item(WeightedItem::unit(value))
807 }
808
809 /// Merge the items of `other` into `self`, leaving `other` empty.
810 ///
811 /// See `.merge_item()` for details.
812 pub fn merge_with(&mut self, other: WeightedList<V,W>) -> &mut Self
813 {
814 for item in other {
815 self.merge_item(item);
816 }
817
818 self
819 }
820
821 /// Merge any duplicate items in the list.
822 ///
823 /// # Usage
824 ///
825 /// ```
826 /// # use weighted_list::*;
827 /// let mut wl = wlist![
828 /// (1, "sup"),
829 /// (2, ""),
830 /// (20, "sup")
831 /// ];
832 ///
833 /// assert_eq!(
834 /// *wl.merge_duplicates(),
835 /// wlist![
836 /// (21, "sup"),
837 /// (2, ""),
838 /// ]
839 /// );
840 /// ```
841 pub fn merge_duplicates(&mut self) -> &mut Self
842 {
843 let orig = std::mem::replace(self, WeightedList::new());
844 self.merge_with(orig);
845 self
846 }
847}
848
849impl<V: Clone, W: Weight> WeightedList<V,W>
850{
851 /// 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.
852 ///
853 /// # Usage
854 ///
855 /// ```
856 /// # use weighted_list::*;
857 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
858 ///
859 /// wl.take_one_at(2);
860 /// assert_eq!( wl, wlist![(2, "sup"), (2, "nova"), (5, "shard")] );
861 ///
862 /// wl.take_one_at(2);
863 /// assert_eq!( wl, wlist![(2, "sup"), (1, "nova"), (5, "shard")] );
864 ///
865 /// wl.take_one_at(2);
866 /// assert_eq!( wl, wlist![(2, "sup"), (5, "shard")] );
867 /// ```
868 pub fn take_one_at(&mut self, weighted_index: W) -> WeightedItem<V,W>
869 {
870 self.take_by_at(W::one(), weighted_index)
871 }
872
873 /// 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.
874 ///
875 /// # Panics
876 ///
877 /// Panics if `weighted_index` is out of bounds.
878 ///
879 /// # Usage
880 ///
881 /// ```
882 /// # use weighted_list::*;
883 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
884 ///
885 /// wl.take_by_at(2, 2);
886 /// assert_eq!( wl, wlist![(2, "sup"), (1, "nova"), (5, "shard")] );
887 ///
888 /// wl.take_by_at(2, 2);
889 /// assert_eq!( wl, wlist![(2, "sup"), (5, "shard")]);
890 /// ```
891 pub fn take_by_at(&mut self, decrement: W, weighted_index: W) -> WeightedItem<V,W>
892 {
893 let idx = self._unweight_index_(weighted_index);
894 let target = &mut self.data[idx];
895
896 target.weight -= decrement;
897
898 if target.weight <= W::zero() {
899 self.data.remove(idx)
900 }
901 else {
902 target.clone()
903 }
904 }
905
906 /// Remove the entire item at `weighted_index`.
907 ///
908 /// # Usage
909 ///
910 /// ```
911 /// # use weighted_list::*;
912 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
913 ///
914 /// wl.take_entire_at(3);
915 /// assert_eq!( wl, wlist![(2, "sup"), (5, "shard")] );
916 /// ```
917 pub fn take_entire_at(&mut self, weighted_index: W) -> WeightedItem<V,W>
918 {
919 self.remove_at(weighted_index)
920 }
921}
922
923// == RANDOMISATION == //
924impl<V, W: Weight> WeightedList<V,W>
925{
926 fn _get_random_weighted_index_up_to_<RNG>(&self, rng: &mut RNG, upper: W) -> Option<W>
927 where RNG: Rng + ?Sized
928 {
929 let len: f64 = nums::cast::<W, f64>(upper)?;
930 let scalar: f64 = rng.random();
931
932 let idx = (len * scalar).floor();
933 let out = nums::cast::<f64, W>(idx)?;
934
935 Some(out)
936 }
937
938 fn _get_random_weighted_index_<RNG>(&self, rng: &mut RNG) -> Option<W>
939 where RNG: Rng + ?Sized
940 {
941 self._get_random_weighted_index_up_to_(rng, self.len())
942 }
943
944 /// Select a random item from the list and return its value, using weighted randomisation.
945 ///
946 /// # Usage
947 ///
948 /// ```
949 /// # use weighted_list::*;
950 /// let wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
951 ///
952 /// wl.select_random_value(&mut rand::rng());
953 /// // could give:
954 /// // - Some("sup" ) with 20% probability
955 /// // - Some("nova" ) with 30% probability
956 /// // - Some("shard") with 50% probability
957 /// ```
958 pub fn select_random_value<RNG>(&self, rng: &mut RNG) -> Option<&V>
959 where RNG: Rng + ?Sized
960 {
961 let out = &self.select_random_item(rng)?.value;
962 Some(out)
963 }
964
965 /// Select a random item from the list, using weighted randomisation.
966 ///
967 /// # Usage
968 ///
969 /// ```
970 /// # use weighted_list::*;
971 /// let wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
972 ///
973 /// wl.select_random_item(&mut rand::rng());
974 /// // could give:
975 /// // - Some(WeightedItem { 2, "sup" }) with 20% probability
976 /// // - Some(WeightedItem { 3, "nova" }) with 30% probability
977 /// // - Some(WeightedItem { 5, "shard" }) with 50% probability
978 /// ```
979 pub fn select_random_item<RNG>(&self, rng: &mut RNG) -> Option<&WeightedItem<V,W>>
980 where RNG: Rng + ?Sized
981 {
982 if self.data.is_empty() { return None }
983
984 let idx = self._get_random_weighted_index_(rng)?;
985 let out = &self[idx];
986
987 Some(out)
988 }
989}
990
991impl<V: Clone, W: Weight> WeightedList<V,W>
992{
993 /// Select a random item from the list using weighted randomisation, and decrement its weight by `1`.
994 ///
995 /// # Usage
996 ///
997 /// ```
998 /// # use weighted_list::*;
999 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
1000 ///
1001 /// wl.take_one_random(&mut rand::rng());
1002 /// // could give:
1003 /// // - Some(WeightedItem { 1, "sup" }) with 20% probability
1004 /// // - Some(WeightedItem { 2, "nova" }) with 30% probability
1005 /// // - Some(WeightedItem { 4, "shard" }) with 50% probability
1006 /// ```
1007 pub fn take_one_random<RNG>(&mut self, rng: &mut RNG) -> Option<WeightedItem<V,W>>
1008 where RNG: Rng + ?Sized
1009 {
1010 self.take_by_random(rng, W::one())
1011 }
1012
1013 /// Select a random item from the list using weighted randomisation, and decrement its weight by `decrement`.
1014 ///
1015 /// # Usage
1016 ///
1017 /// ```
1018 /// # use weighted_list::*;
1019 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
1020 ///
1021 /// wl.take_by_random(&mut rand::rng(), 2);
1022 /// // could give:
1023 /// // - Some(WeightedItem { 0, "sup" }) with 20% probability
1024 /// // - Some(WeightedItem { 1, "nova" }) with 30% probability
1025 /// // - Some(WeightedItem { 3, "shard" }) with 50% probability
1026 /// ```
1027 pub fn take_by_random<RNG>(&mut self, rng: &mut RNG, decrement: W) -> Option<WeightedItem<V,W>>
1028 where RNG: Rng + ?Sized
1029 {
1030 if self.data.is_empty() { return None }
1031
1032 let idx = self._get_random_weighted_index_(rng)?;
1033 let out = self.take_by_at(decrement, idx);
1034
1035 Some(out)
1036 }
1037
1038 /// Select and remove a random item from the list, using weighted randomisation.
1039 ///
1040 /// # Usage
1041 ///
1042 /// ```
1043 /// # use weighted_list::*;
1044 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
1045 ///
1046 /// wl.take_entire_random(&mut rand::rng());
1047 /// // could give:
1048 /// // - Some(WeightedItem { 2, "sup" }) with 20% probability
1049 /// // - Some(WeightedItem { 3, "nova" }) with 30% probability
1050 /// // - Some(WeightedItem { 5, "shard" }) with 50% probability
1051 ///
1052 /// assert!( wl.total_values() == 2 );
1053 /// ```
1054 pub fn take_entire_random<RNG>(&mut self, rng: &mut RNG) -> Option<WeightedItem<V,W>>
1055 where RNG: Rng + ?Sized
1056 {
1057 if self.data.is_empty() { return None }
1058
1059 let idx = self._get_random_weighted_index_(rng)?;
1060 let out = self.take_entire_at(idx);
1061
1062 Some(out)
1063 }
1064}
1065
1066#[bon]
1067impl<V: Clone + Eq, W: Weight> WeightedList<V,W>
1068{
1069 /// Select `count` values using weighted randomisation.
1070 ///
1071 /// Call this method using `bon` builder syntax (see § Usage below).
1072 ///
1073 /// # Options
1074 ///
1075 /// ```ignore
1076 /// rng: RNG,
1077 /// count: usize,
1078 /// replace: Option<bool> = true,
1079 /// decrement: Option<W> = 1,
1080 /// unique: Option<bool> = false,
1081 /// ```
1082 ///
1083 /// - `count`: How many values to select.
1084 /// - `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.
1085 /// - This means at most `self.len()` values will be returned.
1086 /// - `decrement`: How much to decrement weights by if `replace` is `false`.
1087 /// - `unique`: If `true`, only distinct values will be returned.
1088 /// - `replace` becomes irrelevant in this case.
1089 /// - This uses `Eq` equality comparison.
1090 /// - This means at most `self.total_values()` values will be returned.
1091 ///
1092 /// # Usage
1093 ///
1094 /// This method uses the bon builder syntax:
1095 ///
1096 /// ```
1097 /// # use weighted_list::*;
1098 /// let mut pool = wlist![
1099 /// (2, String::from("sup")),
1100 /// (3, String::from("nova")),
1101 /// (5, String::from("shard")),
1102 /// ];
1103 ///
1104 /// let mut rng = rand::rng();
1105 ///
1106 /// // with replacement
1107 /// let selected =
1108 /// pool.select_random_values()
1109 /// .rng(&mut rng)
1110 /// .count(3)
1111 /// .call();
1112 ///
1113 /// assert!(selected.len() == 3);
1114 ///
1115 /// // without replacement
1116 /// let selected =
1117 /// pool.select_random_values()
1118 /// .rng(&mut rng)
1119 /// .count(10)
1120 /// .replace(false)
1121 /// .decrement(2)
1122 /// .call();
1123 ///
1124 /// assert!(selected.len() == 6);
1125 ///
1126 /// // unique only
1127 /// let selected =
1128 /// pool.select_random_values()
1129 /// .rng(&mut rng)
1130 /// .count(100)
1131 /// .unique(true)
1132 /// .call();
1133 ///
1134 /// assert!(selected.len() == 3);
1135 /// ```
1136 ///
1137 /// # Notes
1138 ///
1139 /// - It is not guaranteed that the results will have exactly `count` values.
1140 /// - If `count` exceeds the maximum possible number of values that can be returned, excess iterations will be skipped.
1141 /// - If selection for an iteration fails, that value is excluded from the output list.
1142 /// - This method reserves a `Vec<>` with capacity `count` initially, so be careful of passing in extremely large `count`s.
1143 #[builder]
1144 pub fn select_random_values<RNG>(&self,
1145 rng: &mut RNG,
1146 count: usize,
1147 replace: Option<bool>,
1148 decrement: Option<W>,
1149 unique: Option<bool>,
1150 ) -> Vec<V>
1151 where RNG: Rng + ?Sized
1152 {
1153 let replace = replace.unwrap_or(true);
1154 let decrement = decrement.unwrap_or(W::one());
1155 let unique = unique.unwrap_or(false);
1156
1157 let mut pool = self.clone();
1158 let mut n = 0;
1159 let mut out = Vec::with_capacity(count);
1160
1161 loop
1162 {
1163 n += 1;
1164 if n > count || self.data.is_empty() { break }
1165
1166 if let Some(item) = {
1167 if unique { pool.take_entire_random(rng) }
1168 else if replace { pool.take_by_random(rng, W::zero()) }
1169 else { pool.take_by_random(rng, decrement) }
1170 } {
1171 out.push(item.value.clone());
1172 }
1173 }
1174
1175 out
1176 }
1177
1178 /// Take `count` values using weighted randomisation.
1179 #[builder]
1180 pub fn take_random_values<RNG>(&mut self,
1181 rng: &mut RNG,
1182 count: usize,
1183 take_entire: Option<bool>,
1184 decrement: Option<W>,
1185 ) -> Vec<V>
1186 where RNG: Rng + ?Sized
1187 {
1188 let take_entire = take_entire.unwrap_or(true);
1189 let decrement = decrement.unwrap_or(W::one());
1190
1191 let mut n = 0;
1192 let mut out = Vec::with_capacity(count as usize);
1193
1194 loop
1195 {
1196 n += 1;
1197 if n > count || self.data.is_empty() { break }
1198
1199 if let Some(item) = {
1200 if take_entire { self.take_entire_random(rng) }
1201 else { self.take_by_random(rng, decrement) }
1202 } {
1203 out.push(item.value.clone());
1204 }
1205 }
1206
1207 out
1208 }
1209}
1210
1211#[bon]
1212impl<V: Clone + Eq + std::hash::Hash, W: Weight> WeightedList<V,W>
1213{
1214 /// Variant of `._unweighted_index_()` for mutating random selection with unique outputs.
1215 fn _unweight_index_skipping_(&self,
1216 weighted_index: W,
1217 seen: &std::collections::HashSet<V>,
1218 ) -> Option<usize>
1219 {
1220 let mut t = W::zero();
1221
1222 for (i, item) in self.data.iter().enumerate() {
1223 if seen.contains(&item.value) {
1224 continue
1225 }
1226
1227 t += item.weight;
1228
1229 if t > weighted_index {
1230 return Some(i);
1231 }
1232 }
1233
1234 None
1235 }
1236
1237 /// Take `count` values using weighted randomisation.
1238 #[builder]
1239 pub fn take_random_values_unique<RNG>(&mut self,
1240 rng: &mut RNG,
1241 count: usize,
1242 decrement: Option<W>,
1243 ) -> Vec<V>
1244 where RNG: Rng + ?Sized,
1245 {
1246 let decrement = decrement.unwrap_or(W::one());
1247
1248 let mut n = 0;
1249 let mut l = self.len();
1250
1251 let mut seen = std::collections::HashSet::<V>::new();
1252
1253 let mut out = Vec::with_capacity(
1254 if count > 16 {
1255 count.min(self.total_values())
1256 } else {
1257 count
1258 }
1259 );
1260
1261 loop
1262 {
1263 n += 1;
1264 if n > count || self.data.is_empty() { break }
1265
1266 if let Some(value) = (|| {
1267 let weighted_index = self._get_random_weighted_index_up_to_(rng, l)?;
1268 let idx = self._unweight_index_skipping_(weighted_index, &seen)?;
1269
1270 let target = &mut self.data[idx];
1271 let value = target.value.clone();
1272
1273 target.weight -= decrement;
1274
1275 if target.weight <= W::zero() {
1276 self.data.remove(idx);
1277 }
1278
1279 Some(value)
1280 })()
1281 {
1282 seen.insert(value.clone());
1283 out.push(value.clone());
1284
1285 l = self.data.iter()
1286 .filter_map(
1287 |item| {
1288 if !seen.contains(&item.value) { Some(item.weight) }
1289 else { None }
1290 }
1291 )
1292 .sum::<W>();
1293 }
1294 }
1295
1296 out
1297 }
1298}
1299
1300impl<V: Clone, W: Weight + Clone> WeightedList<V,W>
1301{
1302 /// Shuffle the order of items in the list in-place.
1303 pub fn shuffle_items<RNG>(&mut self, rng: &mut RNG) -> &mut Self
1304 where RNG: Rng + ?Sized
1305 {
1306 self.data.shuffle(rng);
1307 self
1308 }
1309
1310 /// Return a clone with the order of items shuffled.
1311 ///
1312 /// Out-of-place version of `.shuffle_items()`.
1313 pub fn shuffled_items<RNG>(&self, rng: &mut RNG) -> Self
1314 where RNG: Rng + ?Sized
1315 {
1316 let mut out = self.clone();
1317 out.shuffle_items(rng);
1318
1319 out
1320 }
1321
1322 /// Shuffle the pairings of (weight, value) for items in the list, in-place.
1323 ///
1324 /// Values remain in the same order, while weights are re-assigned.
1325 ///
1326 /// # Usage
1327 ///
1328 /// ```
1329 /// # use weighted_list::*;
1330 /// let mut wl = wlist![(2, "sup"), (3, "nova"), (5, "shard")];
1331 ///
1332 /// wl.shuffle_weights(&mut rand::rng());
1333 ///
1334 /// println!("{wl}");
1335 /// // could give:
1336 /// // WeightedList[{3, sup}, {5, nova}, {2, shard}]
1337 ///
1338 /// println!("{wl}");
1339 /// // could now give:
1340 /// // WeightedList[{2, sup}, {5, nova}, {3, shard}]
1341 /// ```
1342 pub fn shuffle_weights<RNG>(&mut self, rng: &mut RNG) -> &mut Self
1343 where RNG: Rng + ?Sized
1344 {
1345 let mut weights: Vec<W> = self.weights().collect();
1346 weights.shuffle(rng);
1347
1348 for item in &mut self.data {
1349 /* guaranteed to be Some */
1350 item.weight = weights.pop().unwrap();
1351 }
1352
1353 self
1354 }
1355
1356 /// Return a clone of the list with (weight, value) pairings shuffled.
1357 ///
1358 /// Out-of-place version of `.shuffle_weights()`.
1359 pub fn shuffled_weights<RNG>(&self, rng: &mut RNG) -> Self
1360 where RNG: Rng + ?Sized
1361 {
1362 let mut out = self.clone();
1363 out.shuffle_weights(rng);
1364
1365 out
1366 }
1367}
1368
1369// == INTERNAL == //
1370impl<V, W: Weight> WeightedList<V,W>
1371{
1372 /// 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`.
1373 fn _unweight_index_nopanic_(&self, weighted_index: W) -> usize
1374 {
1375 let mut t = W::zero();
1376 let mut i = 0;
1377
1378 for item in &self.data {
1379 t += item.weight;
1380
1381 if t > weighted_index {
1382 return i;
1383 }
1384
1385 i += 1;
1386 }
1387
1388 i
1389 }
1390
1391 /// Convert a `weighted_index` to its unweighted equivalent in the underlying `Vec`. Panics on overflow.
1392 fn _unweight_index_(&self, weighted_index: W) -> usize
1393 {
1394 let mut t = W::zero();
1395
1396 for (i, item) in self.data.iter().enumerate() {
1397 t += item.weight;
1398
1399 if t > weighted_index {
1400 return i;
1401 }
1402 }
1403
1404 panic!(
1405 "index out of bounds: the len is {} but the index is {}",
1406 self.len(), weighted_index
1407 );
1408 }
1409}
1410
1411
1412#[cfg(test)] mod tests
1413{
1414 use super::*;
1415
1416 fn el() -> WeightedList<String, i32>
1417 {
1418 wlist![]
1419 }
1420
1421 fn wl() -> WeightedList<String, i32>
1422 {
1423 wlist![
1424 (2, "sup".to_string()),
1425 (3, "nova".to_string()),
1426 (5, "shard".to_string()),
1427 ]
1428 }
1429
1430 #[test] fn _unweight_index_()
1431 {
1432 let list = wl();
1433 assert_eq!( list._unweight_index_(0), 0 );
1434 assert_eq!( list._unweight_index_(1), 0 );
1435 assert_eq!( list._unweight_index_(2), 1 );
1436 assert_eq!( list._unweight_index_(3), 1 );
1437 assert_eq!( list._unweight_index_(4), 1 );
1438 assert_eq!( list._unweight_index_(5), 2 );
1439 assert_eq!( list._unweight_index_(6), 2 );
1440 assert_eq!( list._unweight_index_(7), 2 );
1441 assert_eq!( list._unweight_index_(8), 2 );
1442 assert_eq!( list._unweight_index_(9), 2 );
1443 }
1444
1445 #[test] #[should_panic] fn _unweight_index_empty_()
1446 {
1447 el()._unweight_index_(0);
1448 }
1449
1450 #[test] #[should_panic] fn _unweight_index_out_of_bounds_()
1451 {
1452 wl()._unweight_index_(10);
1453 }
1454
1455 #[test] fn _unweight_index_nopanic_()
1456 {
1457 let list = wl();
1458 assert_eq!( list._unweight_index_nopanic_(10), 3 );
1459 assert_eq!( list._unweight_index_nopanic_(11), 3 );
1460 assert_eq!( list._unweight_index_nopanic_(12), 3 );
1461 }
1462
1463 #[test] fn _unweight_index_skipping_()
1464 {
1465 let list = wl();
1466
1467 let seen = std::collections::HashSet::from(["sup".to_string()]);
1468 assert_eq!( list._unweight_index_skipping_(0, &seen), Some(1) );
1469 assert_eq!( list._unweight_index_skipping_(1, &seen), Some(1) );
1470 assert_eq!( list._unweight_index_skipping_(2, &seen), Some(1) );
1471 assert_eq!( list._unweight_index_skipping_(3, &seen), Some(2) );
1472 assert_eq!( list._unweight_index_skipping_(4, &seen), Some(2) );
1473 assert_eq!( list._unweight_index_skipping_(5, &seen), Some(2) );
1474 assert_eq!( list._unweight_index_skipping_(6, &seen), Some(2) );
1475 assert_eq!( list._unweight_index_skipping_(7, &seen), Some(2) );
1476
1477 let seen = std::collections::HashSet::from(["nova".to_string()]);
1478 assert_eq!( list._unweight_index_skipping_(0, &seen), Some(0) );
1479 assert_eq!( list._unweight_index_skipping_(1, &seen), Some(0) );
1480 assert_eq!( list._unweight_index_skipping_(2, &seen), Some(2) );
1481 assert_eq!( list._unweight_index_skipping_(3, &seen), Some(2) );
1482 assert_eq!( list._unweight_index_skipping_(4, &seen), Some(2) );
1483 assert_eq!( list._unweight_index_skipping_(5, &seen), Some(2) );
1484 assert_eq!( list._unweight_index_skipping_(6, &seen), Some(2) );
1485 }
1486}