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

weighted_list/
weighted_list.rs
