dahl_partition/
lib.rs

1#![allow(dead_code)]
2
3extern crate rand;
4
5use rand::seq::SliceRandom;
6use rand::Rng;
7use std::cmp::Ordering;
8use std::collections::HashMap;
9use std::collections::HashSet;
10use std::convert::TryFrom;
11use std::fmt;
12use std::hash::Hash;
13use std::slice;
14
15/// A partition of `n_items` is a set of subsets such that the subsets are mutually exclusive,
16/// nonempty, and exhaustive.  This data structure enforces mutually exclusivity and provides
17/// methods to test for the other two properties:
18/// [subsets_are_nonempty](struct.Partition.html#method.subsets_are_nonempty) and
19/// [subsets_are_exhaustive](struct.Partition.html#method.subsets_are_exhaustive).
20/// Zero-based numbering is used, e.g., the first
21/// item is index 0 and the first subset is subset 0.  When a subset is added to the data structure,
22/// the next available integer label is used for its subset index.  To remove empty subsets and
23/// relabel subsets into the canonical form, use the
24/// [canonicalize](struct.Partition.html#method.canonicalize) method.
25///
26#[derive(Debug, Clone)]
27pub struct Partition {
28    n_items: usize,
29    n_allocated_items: usize,
30    subsets: Vec<Subset>,
31    labels: Vec<Option<usize>>,
32}
33
34impl Partition {
35    /// Instantiates a partition for `n_items`, none of which are initially allocated.
36    ///
37    /// # Examples
38    ///
39    /// ```
40    /// use dahl_partition::*;
41    /// let partition = Partition::new(10);
42    /// assert_eq!(partition.n_items(), 10);
43    /// assert_eq!(partition.n_subsets(), 0);
44    /// assert_eq!(partition.to_string(), "_ _ _ _ _ _ _ _ _ _");
45    /// ```
46    pub fn new(n_items: usize) -> Self {
47        Self {
48            n_items,
49            n_allocated_items: 0,
50            subsets: Vec::new(),
51            labels: vec![None; n_items],
52        }
53    }
54
55    /// Instantiates a partition for `n_items`, with all items allocated to one subset.
56    ///
57    pub fn one_subset(n_items: usize) -> Self {
58        let mut subset = Subset::new();
59        for i in 0..n_items {
60            subset.add(i);
61        }
62        let labels = vec![Some(0); n_items];
63        Self {
64            n_items,
65            n_allocated_items: n_items,
66            subsets: vec![subset],
67            labels,
68        }
69    }
70
71    /// Instantiates a partition for `n_items`, with each item allocated to its own subset.
72    ///
73    pub fn singleton_subsets(n_items: usize) -> Self {
74        let mut subsets = Vec::with_capacity(n_items);
75        let mut labels = Vec::with_capacity(n_items);
76        for i in 0..n_items {
77            let mut subset = Subset::new();
78            subset.add(i);
79            subsets.push(subset);
80            labels.push(Some(i));
81        }
82        Self {
83            n_items,
84            n_allocated_items: n_items,
85            subsets,
86            labels,
87        }
88    }
89
90    /// Instantiates a partition from a slice of subset labels.
91    /// Items `i` and `j` are in the subset if and only if their labels are equal.
92    ///
93    /// # Examples
94    ///
95    /// ```
96    /// use dahl_partition::*;
97    /// let labels = vec!('a', 'a', 'a', 't', 't', 'a', 'a', 'w', 'a', 'w');
98    /// let partition1 = Partition::from(&labels[..]);
99    /// let partition2 = Partition::from(&[2,2,2,5,5,2,2,7,2,7]);
100    /// let partition3 = Partition::from("AAABBAACAC".as_bytes());
101    /// assert_eq!(partition1, partition2);
102    /// assert_eq!(partition2, partition3);
103    /// assert_eq!(partition1.n_items(), labels.len());
104    /// assert_eq!(partition1.n_subsets(), {
105    ///     let mut x = labels.clone();
106    ///     x.sort();
107    ///     x.dedup();
108    ///     x.len()
109    /// });
110    /// assert_eq!(partition1.to_string(), "0 0 0 1 1 0 0 2 0 2");
111    /// ```
112    pub fn from<T>(labels: &[T]) -> Self
113    where
114        T: Eq + Hash,
115    {
116        let n_items = labels.len();
117        let mut new_labels = vec![None; n_items];
118        let mut subsets = Vec::new();
119        let mut map = HashMap::new();
120        let mut next_label: usize = 0;
121        for i in 0..labels.len() {
122            let key = &labels[i];
123            let label = map.entry(key).or_insert_with(|| {
124                subsets.push(Subset::new());
125                let label = next_label;
126                next_label += 1;
127                label
128            });
129            subsets[*label].add(i);
130            new_labels[i] = Some(*label);
131        }
132        Self {
133            n_items,
134            n_allocated_items: n_items,
135            subsets,
136            labels: new_labels,
137        }
138    }
139
140    /// An iterator over all possible partitions for the specified number of items.
141    ///
142    pub fn iter(n_items: usize) -> PartitionIterator {
143        PartitionIterator {
144            n_items,
145            labels: vec![0; n_items],
146            max: vec![0; n_items],
147            done: false,
148            period: 1,
149        }
150    }
151
152    /// A vector of iterators, which together go over all possible partitions for the specified
153    /// number of items.  This can be useful in multi-thread situations.
154    ///
155    pub fn iter_sharded(n_shards: u32, n_items: usize) -> Vec<PartitionIterator> {
156        let mut shards = Vec::with_capacity(n_shards as usize);
157        let ns = if n_shards == 0 { 1 } else { n_shards };
158        for i in 0..ns {
159            let mut iter = PartitionIterator {
160                n_items,
161                labels: vec![0; n_items],
162                max: vec![0; n_items],
163                done: false,
164                period: ns,
165            };
166            iter.advance(i);
167            shards.push(iter);
168        }
169        shards
170    }
171
172    /// Number of items that can be allocated to the partition.
173    ///
174    /// # Examples
175    ///
176    /// ```
177    /// use dahl_partition::*;
178    /// let partition = Partition::from("AAABBAACAC".as_bytes());
179    /// assert_eq!(partition.n_items(), 10);
180    /// ```
181    pub fn n_items(&self) -> usize {
182        self.n_items
183    }
184
185    /// Number of subsets in the partition.
186    ///
187    /// # Examples
188    ///
189    /// ```
190    /// use dahl_partition::*;
191    /// let mut partition = Partition::from("AAABBAACAC".as_bytes());
192    /// assert_eq!(partition.n_subsets(), 3);
193    /// partition.remove(7);
194    /// partition.remove(9);
195    /// assert_eq!(partition.n_subsets(), 3);
196    /// partition.canonicalize();
197    /// assert_eq!(partition.n_subsets(), 2);
198    /// ```
199    pub fn n_subsets(&self) -> usize {
200        self.subsets.len()
201    }
202
203    /// A vector of length `n_items` whose elements are subset labels.  Panics if any items are
204    /// not allocated.
205    pub fn labels_via_copying(&self) -> Vec<usize> {
206        self.labels.iter().map(|x| x.unwrap()).collect()
207    }
208
209    /// Copy subset labels into a slice.  Panics if any items are not allocated or if the slice
210    /// is not the correct length.
211    pub fn labels_into_slice<T, U>(&self, slice: &mut [T], f: U)
212    where
213        U: Fn(&Option<usize>) -> T,
214    {
215        for (x, y) in slice.iter_mut().zip(self.labels.iter().map(f)) {
216            *x = y;
217        }
218    }
219
220    /// A reference to a vector of length `n_items` whose elements are `None` for items that are
221    /// not allocated (i.e., missing) and, for items that are allocated, `Some(subset_index)` where `subset_index`
222    /// is the index of the subset to which the item is allocated.
223    pub fn labels(&self) -> &Vec<Option<usize>> {
224        &self.labels
225    }
226
227    /// Either `None` for an item that is
228    /// not allocated (i.e., missing) or, for an item that is allocated, `Some(subset_index)` where `subset_index`
229    /// is the index of the subset to which the item is allocated.
230    pub fn label_of(&self, item_index: usize) -> Option<usize> {
231        self.check_item_index(item_index);
232        self.labels[item_index]
233    }
234
235    /// Either `None` for an item that is not allocated (i.e., missing) or, for an item that is allocated,
236    /// `Some(&subset)` where `&subset` is a reference to the subset to which the item is allocated.
237    pub fn subset_of(&self, item_index: usize) -> Option<&Subset> {
238        self.label_of(item_index)
239            .map(|subset_index| &self.subsets[subset_index])
240    }
241
242    /// Returns `true` if and only if both items are allocated (i.e., not missing) and are allocated
243    /// to the same subset.
244    pub fn paired(&self, item1_index: usize, item2_index: usize) -> bool {
245        self.check_item_index(item1_index);
246        self.check_item_index(item2_index);
247        let l1 = self.labels[item1_index];
248        l1.is_some() && l1 == self.labels[item2_index]
249    }
250
251    pub fn subsets(&self) -> &Vec<Subset> {
252        &self.subsets
253    }
254
255    pub fn clean_subset(&mut self, subset_index: usize) {
256        self.check_subset_index(subset_index);
257        self.subsets[subset_index].clean();
258    }
259
260    /// Add a new empty subset to the partition.
261    ///
262    /// # Examples
263    ///
264    /// ```
265    /// use dahl_partition::*;
266    /// let mut partition = Partition::from("AAABBAACAC".as_bytes());
267    /// assert_eq!(partition.n_subsets(), 3);
268    /// assert!(partition.subsets_are_nonempty());
269    /// partition.new_subset();
270    /// assert!(!partition.subsets_are_nonempty());
271    /// assert_eq!(partition.n_subsets(), 4);
272    /// ```
273    pub fn new_subset(&mut self) {
274        self.subsets.push(Subset::new());
275    }
276
277    /// Add a new subset containing `item_index` to the partition.
278    ///
279    /// # Examples
280    ///
281    /// ```
282    /// use dahl_partition::*;
283    /// let mut partition = Partition::new(3);
284    /// partition.add(1);
285    /// assert_eq!(partition.to_string(), "_ 0 _");
286    /// partition.add(0);
287    /// assert_eq!(partition.to_string(), "1 0 _");
288    /// partition.canonicalize();
289    /// assert_eq!(partition.to_string(), "0 1 _");
290    /// ```
291    pub fn add(&mut self, item_index: usize) -> &mut Self {
292        self.check_item_index(item_index);
293        self.check_not_allocated(item_index);
294        self.n_allocated_items += 1;
295        self.subsets.push(Subset::new());
296        self.add_engine(item_index, self.subsets.len() - 1);
297        self
298    }
299
300    /// Add item `item_index` to the subset `subset_index` of the partition.
301    ///
302    /// # Examples
303    ///
304    /// ```
305    /// use dahl_partition::*;
306    /// let mut partition = Partition::new(3);
307    /// partition.add(1);
308    /// assert_eq!(partition.to_string(), "_ 0 _");
309    /// partition.add_with_index(2, 0);
310    /// assert_eq!(partition.to_string(), "_ 0 0");
311    /// ```
312    pub fn add_with_index(&mut self, item_index: usize, subset_index: usize) -> &mut Self {
313        self.check_item_index(item_index);
314        self.check_not_allocated(item_index);
315        self.check_subset_index(subset_index);
316        self.n_allocated_items += 1;
317        self.add_engine(item_index, subset_index);
318        self
319    }
320
321    /// Remove item `item_index` from the partition.
322    ///
323    /// # Examples
324    ///
325    /// ```
326    /// use dahl_partition::*;
327    /// let mut partition = Partition::from("AAABBAACAC".as_bytes());
328    /// partition.remove(1);
329    /// partition.remove(4);
330    /// assert_eq!(partition.to_string(), "0 _ 0 1 _ 0 0 2 0 2");
331    /// partition.remove(3);
332    /// assert_eq!(partition.to_string(), "0 _ 0 _ _ 0 0 2 0 2");
333    /// partition.canonicalize();
334    /// assert_eq!(partition.to_string(), "0 _ 0 _ _ 0 0 1 0 1");
335    /// ```
336    pub fn remove(&mut self, item_index: usize) -> &mut Self {
337        self.check_item_index(item_index);
338        self.remove_engine(item_index, self.check_allocated(item_index));
339        self
340    }
341
342    /// Remove item `item_index` from the partition and, if this creates an empty subset, relabel.
343    ///
344    /// # Examples
345    ///
346    /// ```
347    /// use dahl_partition::*;
348    /// let mut partition = Partition::from("AAABBAACAC".as_bytes());
349    /// partition.remove_clean_and_relabel(1, |x,y| {});
350    /// partition.remove_clean_and_relabel(4, |x,y| {});
351    /// partition.remove_clean_and_relabel(3, |x,y| {});
352    /// assert_eq!(partition.to_string(), "0 _ 0 _ _ 0 0 1 0 1");
353    /// ```
354    pub fn remove_clean_and_relabel<T>(&mut self, item_index: usize, mut callback: T) -> &mut Self
355    where
356        T: FnMut(usize, usize),
357    {
358        self.check_item_index(item_index);
359        let subset_index = self.check_allocated(item_index);
360        self.remove_engine(item_index, subset_index);
361        if self.subsets[subset_index].is_empty() {
362            let moved_subset_index = self.subsets.len() - 1;
363            if moved_subset_index != subset_index {
364                for i in self.subsets[moved_subset_index].items() {
365                    self.labels[*i] = Some(subset_index);
366                }
367            }
368            callback(subset_index, moved_subset_index);
369            self.clean_subset(subset_index);
370            self.subsets.swap_remove(subset_index);
371        } else {
372            self.subsets[subset_index].clean();
373        }
374        self
375    }
376
377    /// Remove item `item_index` from the subset `subset_index` of the partition.
378    ///
379    /// # Examples
380    ///
381    /// ```
382    /// use dahl_partition::*;
383    /// let mut partition = Partition::from("AAABBAACAC".as_bytes());
384    /// partition.remove_with_index(1,0);
385    /// partition.remove_with_index(4,1);
386    /// assert_eq!(partition.to_string(), "0 _ 0 1 _ 0 0 2 0 2");
387    /// partition.remove_with_index(3,1);
388    /// assert_eq!(partition.to_string(), "0 _ 0 _ _ 0 0 2 0 2");
389    /// partition.canonicalize();
390    /// assert_eq!(partition.to_string(), "0 _ 0 _ _ 0 0 1 0 1");
391    /// ```
392    pub fn remove_with_index(&mut self, item_index: usize, subset_index: usize) -> &mut Self {
393        self.check_item_index(item_index);
394        self.check_item_in_subset(item_index, subset_index);
395        self.remove_engine(item_index, subset_index);
396        self
397    }
398
399    /// Removes that last item from the subset `subset_index` of the partition.  This may or may not
400    /// be the most recently added item.
401    ///
402    /// # Examples
403    ///
404    /// ```
405    /// use dahl_partition::*;
406    /// let mut partition = Partition::from("AAABBAACAC".as_bytes());
407    /// assert_eq!(partition.pop_item(1),Some(4));
408    /// ```
409    pub fn pop_item(&mut self, subset_index: usize) -> Option<usize> {
410        self.check_subset_index(subset_index);
411        let i_option = self.subsets[subset_index].pop();
412        if let Some(item_index) = i_option {
413            self.labels[item_index] = None;
414            self.n_allocated_items -= 1;
415        }
416        i_option
417    }
418
419    /// Removes that last subset of the partition if it empty.
420    ///
421    /// # Examples
422    ///
423    /// ```
424    /// use dahl_partition::*;
425    /// let mut partition = Partition::from("AAABBAACAC".as_bytes());
426    /// partition.new_subset();
427    /// assert_ne!(partition.pop_subset(),None);
428    /// ```
429    pub fn pop_subset(&mut self) -> Option<Subset> {
430        if self.subsets.is_empty() {
431            return None;
432        }
433        let subset = self.subsets.last().unwrap();
434        if !subset.is_empty() {
435            return None;
436        }
437        self.subsets.pop()
438    }
439
440    /// Transfers item `item_index` to a new empty subset.
441    pub fn transfer(&mut self, item_index: usize) -> &mut Self {
442        self.check_item_index(item_index);
443        let subset_index = self.check_allocated(item_index);
444        self.subsets[subset_index].remove(item_index);
445        self.subsets.push(Subset::new());
446        self.add_engine(item_index, self.subsets.len() - 1);
447        self
448    }
449
450    /// Transfers item `item_index` from the subset `old_subset_index` to another subset
451    /// `new_subset_index`.
452    pub fn transfer_with_index(
453        &mut self,
454        item_index: usize,
455        old_subset_index: usize,
456        new_subset_index: usize,
457    ) -> &mut Self {
458        self.check_item_index(item_index);
459        self.check_item_in_subset(item_index, old_subset_index);
460        self.check_subset_index(new_subset_index);
461        self.subsets[old_subset_index].remove(item_index);
462        self.add_engine(item_index, new_subset_index);
463        self
464    }
465
466    /// Put partition into canonical form, removing empty subsets and consecutively numbering
467    /// subsets starting at 0.
468    ///
469    /// # Examples
470    ///
471    /// ```
472    /// use dahl_partition::*;
473    /// let mut partition = Partition::new(3);
474    /// partition.add(1);
475    /// partition.add(2);
476    /// partition.add(0);
477    /// assert_eq!(partition.to_string(), "2 0 1");
478    /// partition.canonicalize();
479    /// assert_eq!(partition.to_string(), "0 1 2");
480    /// ```
481    ///
482    pub fn canonicalize(&mut self) -> &mut Self {
483        self.canonicalize_by_permutation(None)
484    }
485
486    /// Put partition into canonical form according to a permutation, removing empty subsets and consecutively numbering
487    /// subsets starting at 0 based on the supplied permutation.  If `None` is supplied, the natural permutation is used.
488    ///
489    /// # Examples
490    ///
491    /// ```
492    /// use dahl_partition::*;
493    /// use std::fs::canonicalize;
494    /// let mut partition = Partition::new(3);
495    /// partition.add(1);
496    /// partition.add(2);
497    /// partition.add(0);
498    /// assert_eq!(partition.to_string(), "2 0 1");
499    /// partition.canonicalize_by_permutation(Some(&Permutation::from_slice(&[0,2,1]).unwrap()));
500    /// assert_eq!(partition.to_string(), "0 2 1");
501    /// ```
502    ///
503    pub fn canonicalize_by_permutation(&mut self, permutation: Option<&Permutation>) -> &mut Self {
504        if let Some(p) = permutation {
505            assert_eq!(self.n_items, p.len());
506        };
507        if self.n_allocated_items == 0 {
508            return self;
509        }
510        let mut new_labels = vec![None; self.n_items];
511        let mut next_label: usize = 0;
512        for i in 0..self.n_items {
513            let ii = match permutation {
514                None => i,
515                Some(p) => p[i],
516            };
517            match new_labels[ii] {
518                Some(_) => (),
519                None => match self.labels[ii] {
520                    None => (),
521                    Some(subset_index) => {
522                        let subset = &mut self.subsets[subset_index];
523                        subset.clean();
524                        let some_label = Some(next_label);
525                        next_label += 1;
526                        for j in subset.vector.iter() {
527                            new_labels[*j] = some_label;
528                        }
529                    }
530                },
531            }
532        }
533        self.subsets.sort_unstable_by(|x, y| {
534            if x.is_empty() {
535                if y.is_empty() {
536                    Ordering::Equal
537                } else {
538                    Ordering::Greater
539                }
540            } else if y.is_empty() {
541                Ordering::Less
542            } else {
543                new_labels[x.vector[0]]
544                    .unwrap()
545                    .cmp(&new_labels[y.vector[0]].unwrap())
546            }
547        });
548        while !self.subsets.is_empty() && self.subsets.last().unwrap().is_empty() {
549            self.subsets.pop();
550        }
551        self.labels = new_labels;
552        self
553    }
554
555    /// Test whether subsets are exhaustive.
556    pub fn subsets_are_exhaustive(&self) -> bool {
557        self.n_allocated_items == self.n_items
558    }
559
560    /// Test whether subsets are nonempty.
561    pub fn subsets_are_nonempty(&self) -> bool {
562        self.subsets.iter().all(|x| !x.is_empty())
563    }
564
565    fn add_engine(&mut self, item_index: usize, subset_index: usize) {
566        self.labels[item_index] = Some(subset_index);
567        self.subsets[subset_index].add(item_index);
568    }
569
570    fn remove_engine(&mut self, item_index: usize, subset_index: usize) {
571        self.labels[item_index] = None;
572        self.subsets[subset_index].remove(item_index);
573        self.n_allocated_items -= 1;
574    }
575
576    fn check_item_index(&self, item_index: usize) {
577        if item_index >= self.n_items {
578            panic!(
579                "Attempt to allocate item {} when only {} are available.",
580                item_index, self.n_items
581            );
582        };
583    }
584
585    fn check_subset_index(&self, subset_index: usize) {
586        if subset_index >= self.n_subsets() {
587            panic!(
588                "Attempt to allocate to subset {} when only {} are available.",
589                subset_index,
590                self.n_subsets()
591            );
592        };
593    }
594
595    fn check_allocated(&self, item_index: usize) -> usize {
596        match self.labels[item_index] {
597            Some(subset_index) => subset_index,
598            None => panic!("Item {} is not allocated.", item_index),
599        }
600    }
601
602    fn check_not_allocated(&self, item_index: usize) {
603        if let Some(j) = self.labels[item_index] {
604            panic!("Item {} is already allocated to subset {}.", item_index, j)
605        };
606    }
607
608    fn check_item_in_subset(&self, item_index: usize, subset_index: usize) {
609        match self.labels[item_index] {
610            Some(j) => {
611                if j != subset_index {
612                    panic!("Item {} is already allocated to subset {}.", item_index, j);
613                };
614            }
615            None => panic!("Item {} is not allocated to any subset.", item_index),
616        };
617    }
618}
619
620impl fmt::Display for Partition {
621    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
622        let mut str = "";
623        for i in 0..self.n_items {
624            fmt.write_str(str)?;
625            str = " ";
626            match self.labels[i] {
627                None => fmt.write_str("_")?,
628                Some(subset_index) => {
629                    let x = subset_index.to_string();
630                    fmt.write_str(&x[..])?;
631                }
632            };
633        }
634        Ok(())
635    }
636}
637
638impl PartialEq for Partition {
639    fn eq(&self, other: &Self) -> bool {
640        self.n_items() == other.n_items
641            && self.n_allocated_items == other.n_allocated_items
642            && self.subsets.len() == other.subsets.len()
643            && {
644                let mut x = self.clone();
645                let x1 = x.canonicalize();
646                let mut y = other.clone();
647                let y1 = y.canonicalize();
648                x1.subsets == y1.subsets
649            }
650    }
651}
652
653#[cfg(test)]
654mod tests_partition {
655    use super::*;
656
657    #[test]
658    fn test_complete() {
659        let mut p = Partition::new(6);
660        p.add(2);
661        for i in &[3, 4, 5] {
662            p.add_with_index(*i, 0);
663        }
664        assert_eq!(p.to_string(), "_ _ 0 0 0 0");
665        assert_eq!(p.subsets_are_exhaustive(), false);
666        p.add(1);
667        p.add_with_index(0, 1);
668        p.canonicalize();
669        assert_eq!(p.to_string(), "0 0 1 1 1 1");
670        p.remove_with_index(2, 1);
671        p.canonicalize();
672        assert_eq!(p.to_string(), "0 0 _ 1 1 1");
673        assert_ne!(p.subsets_are_exhaustive(), true);
674        p.remove_with_index(0, 0);
675        p.remove_with_index(1, 0);
676        assert_eq!(p.to_string(), "_ _ _ 1 1 1");
677        p.canonicalize();
678        assert_eq!(p.to_string(), "_ _ _ 0 0 0");
679        p.transfer(5);
680        p.transfer_with_index(3, 0, 1);
681        assert_eq!(p.to_string(), "_ _ _ 1 0 1");
682        p.add_with_index(1, 1);
683        assert_eq!(p.to_string(), "_ 1 _ 1 0 1");
684        p.canonicalize();
685        assert_eq!(p.to_string(), "_ 0 _ 0 1 0");
686        assert_eq!(p.n_subsets(), 2);
687        p.remove_with_index(4, 1);
688        assert_eq!(p.n_subsets(), 2);
689        p.canonicalize();
690        assert_eq!(p.n_subsets(), 1);
691    }
692
693    #[test]
694    fn test_canonicalize() {
695        let mut p = Partition::new(6);
696        p.add(0)
697            .add_with_index(1, 0)
698            .add(2)
699            .add_with_index(3, 1)
700            .add_with_index(4, 0);
701        assert_eq!(p.to_string(), "0 0 1 1 0 _");
702        p.add(5);
703        assert_eq!(p.to_string(), "0 0 1 1 0 2");
704        p.remove(2);
705        p.remove(3);
706        assert_eq!(p.to_string(), "0 0 _ _ 0 2");
707        p.canonicalize();
708        assert_eq!(p.to_string(), "0 0 _ _ 0 1");
709    }
710
711    #[test]
712    fn test_from() {
713        let p0 = Partition::from(&[45, 34, 23, 23, 23, 24, 45]);
714        assert_eq!(p0.to_string(), "0 1 2 2 2 3 0");
715        let p1 = Partition::from("ABCAADAABD".as_bytes());
716        assert_eq!(p1.to_string(), "0 1 2 0 0 3 0 0 1 3");
717        let p2 = Partition::from(p1.labels());
718        assert_eq!(p1.to_string(), p2.to_string());
719        let p3 = Partition::from(p0.labels());
720        assert_eq!(p0.to_string(), p3.to_string());
721    }
722
723    #[test]
724    fn test_equality() {
725        let p0 = Partition::from("ABCAADAABD".as_bytes());
726        let p1 = Partition::from("ABCAADAABD".as_bytes());
727        let p2 = Partition::from("ABCRRRRRRD".as_bytes());
728        assert_eq!(p0, p1);
729        assert_ne!(p0, p2);
730    }
731}
732
733#[doc(hidden)]
734pub struct PartitionIterator {
735    n_items: usize,
736    labels: Vec<usize>,
737    max: Vec<usize>,
738    done: bool,
739    period: u32,
740}
741
742impl PartitionIterator {
743    fn advance(&mut self, times: u32) {
744        for _ in 0..times {
745            let mut i = self.n_items - 1;
746            while (i > 0) && (self.labels[i] == self.max[i - 1] + 1) {
747                self.labels[i] = 0;
748                self.max[i] = self.max[i - 1];
749                i -= 1;
750            }
751            if i == 0 {
752                self.done = true;
753                return;
754            }
755            self.labels[i] += 1;
756            let m = self.max[i].max(self.labels[i]);
757            self.max[i] = m;
758            i += 1;
759            while i < self.n_items {
760                self.max[i] = m;
761                self.labels[i] = 0;
762                i += 1;
763            }
764        }
765    }
766}
767
768impl Iterator for PartitionIterator {
769    type Item = Vec<usize>;
770
771    fn next(&mut self) -> Option<Self::Item> {
772        if self.done {
773            None
774        } else {
775            let result = Some(self.labels.clone());
776            self.advance(self.period);
777            result
778        }
779    }
780}
781
782#[cfg(test)]
783mod tests_iterator {
784    use super::*;
785
786    fn test_iterator_engine(n_items: usize, bell_number: usize) {
787        let mut counter = 0usize;
788        for _ in Partition::iter(n_items) {
789            counter += 1;
790        }
791        assert_eq!(counter, bell_number);
792    }
793
794    #[test]
795    fn test_iterator() {
796        test_iterator_engine(1, 1);
797        test_iterator_engine(2, 2);
798        let mut ph = PartitionsHolder::with_capacity(5, 3);
799        for labels in Partition::iter(3) {
800            ph.push_slice(&labels[..]);
801        }
802        assert_eq!(
803            ph.view().data(),
804            &[0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1, 0, 1, 2]
805        );
806        test_iterator_engine(10, 115_975);
807    }
808
809    #[test]
810    fn test_sharded() {
811        let mut ph = PartitionsHolder::new(3);
812        for iter in Partition::iter_sharded(8, 3) {
813            for labels in iter {
814                ph.push_slice(&labels[..]);
815            }
816        }
817        assert_eq!(
818            ph.view().data(),
819            &[0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1, 0, 1, 2]
820        );
821    }
822}
823
824/// A data structure representing sets of integers.  The user instantiates subsets through
825/// [Partition](struct.Partition.html).
826///
827#[derive(Debug, Clone)]
828pub struct Subset {
829    n_items: usize,
830    set: HashSet<usize>,
831    vector: Vec<usize>,
832    is_clean: bool,
833}
834
835impl Subset {
836    pub fn new() -> Subset {
837        Subset {
838            n_items: 0,
839            set: HashSet::new(),
840            vector: Vec::new(),
841            is_clean: true,
842        }
843    }
844
845    fn from<'a, I>(x: I) -> Subset
846    where
847        I: Iterator<Item = &'a usize>,
848    {
849        let mut set = HashSet::new();
850        let mut vector = Vec::new();
851        let mut n_items = 0;
852        for i in x {
853            if set.insert(*i) {
854                n_items += 1;
855                vector.push(*i);
856            }
857        }
858        Subset {
859            n_items,
860            set,
861            vector,
862            is_clean: true,
863        }
864    }
865
866    pub fn add(&mut self, i: usize) -> bool {
867        if self.set.insert(i) {
868            self.n_items += 1;
869            if self.is_clean {
870                self.vector.push(i);
871            }
872            true
873        } else {
874            false
875        }
876    }
877
878    fn merge(&mut self, other: &Subset) -> bool {
879        let mut added = false;
880        if other.is_clean {
881            for i in other.vector.iter() {
882                added = self.add(*i) || added;
883            }
884        } else {
885            for i in other.set.iter() {
886                added = self.add(*i) || added;
887            }
888        }
889        added
890    }
891
892    pub fn intersection_count(&self, other: &Subset) -> usize {
893        if self.n_items > other.n_items {
894            other.intersection_count(self)
895        } else {
896            let mut count = 0;
897            if self.is_clean {
898                for i in self.vector.iter() {
899                    if other.contains(*i) {
900                        count += 1;
901                    };
902                }
903            } else if other.is_clean {
904                for i in other.vector.iter() {
905                    if self.contains(*i) {
906                        count += 1;
907                    };
908                }
909            } else {
910                for i in self.set.iter() {
911                    if other.contains(*i) {
912                        count += 1;
913                    };
914                }
915            };
916            count
917        }
918    }
919
920    pub fn intersection(&self, other: &Subset) -> Subset {
921        let set: HashSet<usize> = self.set.intersection(&other.set).copied().collect();
922        Subset {
923            n_items: set.len(),
924            set,
925            vector: Vec::new(),
926            is_clean: false,
927        }
928    }
929
930    pub fn contains(&self, i: usize) -> bool {
931        self.set.contains(&i)
932    }
933
934    fn remove(&mut self, i: usize) -> bool {
935        if self.set.remove(&i) {
936            self.n_items -= 1;
937            self.vector.clear();
938            self.is_clean = false;
939            true
940        } else {
941            false
942        }
943    }
944
945    fn pop(&mut self) -> Option<usize> {
946        if !self.is_clean {
947            self.clean()
948        }
949        let i_option = self.vector.pop();
950        if let Some(i) = i_option {
951            self.set.remove(&i);
952            self.n_items -= 1;
953        };
954        i_option
955    }
956
957    fn clean(&mut self) {
958        if !self.is_clean {
959            for i in &self.set {
960                self.vector.push(*i);
961            }
962            self.is_clean = true;
963        }
964    }
965
966    /// The number of items in the subset.
967    pub fn n_items(&self) -> usize {
968        self.n_items
969    }
970
971    /// Is the subset empty?
972    pub fn is_empty(&self) -> bool {
973        self.n_items == 0
974    }
975
976    /// A reference to the elements of the set.
977    pub fn items(&self) -> &Vec<usize> {
978        if !self.is_clean {
979            panic!("Subset is not clean.  Please clean it first.");
980        }
981        &self.vector
982    }
983
984    /// A copy of the elements of the set.
985    pub fn items_via_copying(&self) -> Vec<usize> {
986        if self.is_clean {
987            self.vector.clone()
988        } else {
989            let mut vector = Vec::new();
990            for i in &self.set {
991                vector.push(*i);
992            }
993            vector
994        }
995    }
996}
997
998impl fmt::Display for Subset {
999    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
1000        let mut clone = self.clone();
1001        clone.clean();
1002        clone.vector.sort_unstable();
1003        fmt.write_str("{")?;
1004        let mut str = "";
1005        for i in clone.vector {
1006            fmt.write_str(str)?;
1007            fmt.write_str(&format!("{}", i))?;
1008            str = ",";
1009        }
1010        fmt.write_str("}")?;
1011        Ok(())
1012    }
1013}
1014
1015impl Default for Subset {
1016    fn default() -> Self {
1017        Self::new()
1018    }
1019}
1020
1021impl PartialEq for Subset {
1022    fn eq(&self, other: &Self) -> bool {
1023        self.set == other.set
1024    }
1025}
1026
1027#[cfg(test)]
1028mod tests_subset {
1029    use super::*;
1030
1031    #[test]
1032    fn test_add() {
1033        let mut s = Subset::new();
1034        assert_eq!(s.add(0), true);
1035        assert_eq!(s.add(1), true);
1036        assert_eq!(s.add(0), false);
1037        assert_eq!(s.add(1), false);
1038        assert_eq!(s.add(0), false);
1039        assert_eq!(s.n_items(), 2);
1040        assert_eq!(s.to_string(), "{0,1}");
1041    }
1042
1043    #[test]
1044    fn test_merge() {
1045        let mut s1 = Subset::from([2, 1, 6, 2].iter());
1046        let mut s2 = Subset::from([0, 2, 1].iter());
1047        s1.merge(&mut s2);
1048        assert_eq!(s1.to_string(), "{0,1,2,6}");
1049        let mut s3 = Subset::from([7, 2].iter());
1050        s3.merge(&mut s1);
1051        assert_eq!(s3.to_string(), "{0,1,2,6,7}");
1052    }
1053
1054    #[test]
1055    fn test_remove() {
1056        let mut s = Subset::from([2, 1, 6, 2].iter());
1057        assert_eq!(s.remove(0), false);
1058        assert_eq!(s.remove(2), true);
1059        assert_eq!(s.to_string(), "{1,6}");
1060    }
1061
1062    #[test]
1063    fn test_equality() {
1064        let s1 = Subset::from([2, 1, 6, 2].iter());
1065        let s2 = Subset::from([6, 1, 2].iter());
1066        assert_eq!(s1, s2);
1067        let s3 = Subset::from([6, 1, 2, 3].iter());
1068        assert_ne!(s1, s3);
1069    }
1070}
1071
1072/// A data structure holding partitions.
1073///
1074pub struct PartitionsHolder {
1075    data: Vec<i32>,
1076    n_partitions: usize,
1077    n_items: usize,
1078    by_row: bool,
1079}
1080
1081impl PartitionsHolder {
1082    pub fn new(n_items: usize) -> PartitionsHolder {
1083        PartitionsHolder {
1084            data: Vec::new(),
1085            n_partitions: 0,
1086            n_items,
1087            by_row: false,
1088        }
1089    }
1090
1091    pub fn with_capacity(capacity: usize, n_items: usize) -> PartitionsHolder {
1092        PartitionsHolder {
1093            data: Vec::with_capacity(capacity * n_items),
1094            n_partitions: 0,
1095            n_items,
1096            by_row: false,
1097        }
1098    }
1099
1100    pub fn allocated(n_partitions: usize, n_items: usize, by_row: bool) -> PartitionsHolder {
1101        PartitionsHolder {
1102            data: vec![0; n_partitions * n_items],
1103            n_partitions,
1104            n_items,
1105            by_row,
1106        }
1107    }
1108
1109    pub fn enumerated(n_items: usize) -> PartitionsHolder {
1110        let mut ph = PartitionsHolder::new(n_items);
1111        for partition in Partition::iter(n_items) {
1112            ph.push_slice(&partition[..]);
1113        }
1114        ph
1115    }
1116
1117    pub fn n_partitions(&self) -> usize {
1118        self.n_partitions
1119    }
1120
1121    pub fn n_items(&self) -> usize {
1122        self.n_items
1123    }
1124
1125    pub fn by_row(&self) -> bool {
1126        self.by_row
1127    }
1128
1129    pub fn push_slice(&mut self, partition: &[usize]) {
1130        assert!(!self.by_row, "Pushing requires that by_row = false.");
1131        assert_eq!(
1132            partition.len(),
1133            self.n_items,
1134            "Inconsistent number of items."
1135        );
1136        for j in partition {
1137            self.data.push(i32::try_from(*j).unwrap())
1138        }
1139        self.n_partitions += 1
1140    }
1141
1142    pub fn push_partition(&mut self, partition: &Partition) {
1143        assert!(!self.by_row, "Pushing requires that by_row = false.");
1144        assert_eq!(partition.n_items, self.n_items);
1145        for j in partition.labels() {
1146            self.data.push(i32::try_from(j.unwrap()).unwrap())
1147        }
1148        self.n_partitions += 1
1149    }
1150
1151    pub fn view(&mut self) -> PartitionsHolderBorrower {
1152        PartitionsHolderBorrower::from_slice(
1153            &mut self.data[..],
1154            self.n_partitions,
1155            self.n_items,
1156            self.by_row,
1157        )
1158    }
1159}
1160
1161#[doc(hidden)]
1162pub struct PartitionsHolderBorrower<'a> {
1163    data: &'a mut [i32],
1164    n_partitions: usize,
1165    n_items: usize,
1166    by_row: bool,
1167    index: usize,
1168}
1169
1170impl std::ops::Index<(usize, usize)> for PartitionsHolderBorrower<'_> {
1171    type Output = i32;
1172    fn index(&self, (i, j): (usize, usize)) -> &Self::Output {
1173        if self.by_row {
1174            &self.data[self.n_partitions * j + i]
1175        } else {
1176            &self.data[self.n_items * i + j]
1177        }
1178    }
1179}
1180
1181impl std::ops::IndexMut<(usize, usize)> for PartitionsHolderBorrower<'_> {
1182    fn index_mut(&mut self, (i, j): (usize, usize)) -> &mut Self::Output {
1183        if self.by_row {
1184            &mut self.data[self.n_partitions * j + i]
1185        } else {
1186            &mut self.data[self.n_items * i + j]
1187        }
1188    }
1189}
1190
1191impl<'a> PartitionsHolderBorrower<'a> {
1192    pub fn from_slice(
1193        data: &'a mut [i32],
1194        n_partitions: usize,
1195        n_items: usize,
1196        by_row: bool,
1197    ) -> Self {
1198        assert_eq!(data.len(), n_partitions * n_items);
1199        Self {
1200            data,
1201            n_partitions,
1202            n_items,
1203            by_row,
1204            index: 0,
1205        }
1206    }
1207
1208    /// # Safety
1209    ///
1210    /// Added for FFI.
1211    pub unsafe fn from_ptr(
1212        data: *mut i32,
1213        n_partitions: usize,
1214        n_items: usize,
1215        by_row: bool,
1216    ) -> Self {
1217        let data = slice::from_raw_parts_mut(data, n_partitions * n_items);
1218        Self {
1219            data,
1220            n_partitions,
1221            n_items,
1222            by_row,
1223            index: 0,
1224        }
1225    }
1226    pub fn n_partitions(&self) -> usize {
1227        self.n_partitions
1228    }
1229
1230    pub fn n_items(&self) -> usize {
1231        self.n_items
1232    }
1233
1234    pub fn by_row(&self) -> bool {
1235        self.by_row
1236    }
1237
1238    pub fn data(&self) -> &[i32] {
1239        self.data
1240    }
1241
1242    /// # Safety
1243    ///
1244    /// You're on your own with the indices.
1245    pub unsafe fn get_unchecked(&self, (i, j): (usize, usize)) -> &i32 {
1246        if self.by_row {
1247            self.data.get_unchecked(self.n_partitions * j + i)
1248        } else {
1249            self.data.get_unchecked(self.n_items * i + j)
1250        }
1251    }
1252
1253    /// # Safety
1254    ///
1255    /// You're on your own with the indices.
1256    pub unsafe fn get_unchecked_mut(&mut self, (i, j): (usize, usize)) -> &mut i32 {
1257        if self.by_row {
1258            self.data.get_unchecked_mut(self.n_partitions * j + i)
1259        } else {
1260            self.data.get_unchecked_mut(self.n_items * i + j)
1261        }
1262    }
1263
1264    pub fn get(&self, i: usize) -> Partition {
1265        if self.by_row {
1266            let mut labels = Vec::with_capacity(self.n_items);
1267            for j in 0..self.n_items {
1268                labels.push(self.data[self.n_partitions * j + i]);
1269            }
1270            Partition::from(&labels[..])
1271        } else {
1272            Partition::from(&self.data[(i * self.n_items)..((i + 1) * self.n_items)])
1273        }
1274    }
1275
1276    pub fn get_all(&self) -> Vec<Partition> {
1277        let mut x = Vec::with_capacity(self.n_partitions);
1278        for k in 0..self.n_partitions {
1279            x.push(self.get(k))
1280        }
1281        x
1282    }
1283
1284    pub fn push_slice(&mut self, partition: &[usize]) {
1285        assert_eq!(
1286            partition.len(),
1287            self.n_items,
1288            "Inconsistent number of items."
1289        );
1290        for (j, v) in partition.iter().enumerate() {
1291            let v = i32::try_from(*v).unwrap();
1292            let o = if self.by_row {
1293                self.n_partitions * j + self.index
1294            } else {
1295                self.n_items * self.index + j
1296            };
1297            unsafe {
1298                *self.data.get_unchecked_mut(o) = v;
1299            }
1300        }
1301        self.index += 1;
1302    }
1303
1304    pub fn push_partition(&mut self, partition: &Partition) {
1305        assert!(
1306            self.index < self.n_partitions,
1307            "The holder has capacity {} so cannot push with index {}.",
1308            self.n_partitions,
1309            self.index
1310        );
1311        assert_eq!(
1312            partition.n_items(),
1313            self.n_items,
1314            "Inconsistent number of items."
1315        );
1316        for (j, v) in partition.labels().iter().enumerate() {
1317            let v = i32::try_from(v.unwrap()).unwrap();
1318            let o = if self.by_row {
1319                self.n_partitions * j + self.index
1320            } else {
1321                self.n_items * self.index + j
1322            };
1323            unsafe {
1324                *self.data.get_unchecked_mut(o) = v;
1325            }
1326        }
1327        self.index += 1;
1328    }
1329}
1330
1331impl<'a> fmt::Display for PartitionsHolderBorrower<'a> {
1332    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
1333        for i in 0..self.n_partitions {
1334            let x = self.get(i).to_string();
1335            fmt.write_str(&x[..])?;
1336            fmt.write_str("\n")?;
1337        }
1338        Ok(())
1339    }
1340}
1341
1342#[cfg(test)]
1343mod tests_partitions_holder {
1344    use super::*;
1345
1346    #[test]
1347    fn test_partition_holder() {
1348        let mut phf = PartitionsHolder::new(4);
1349        assert!(!phf.by_row());
1350        assert_eq!(phf.n_partitions(), 0);
1351        phf.push_partition(&Partition::from("AABC".as_bytes()));
1352        phf.push_partition(&Partition::from("ABCD".as_bytes()));
1353        phf.push_partition(&Partition::from("AABA".as_bytes()));
1354        assert_eq!(phf.n_partitions(), 3);
1355        let phfv = phf.view();
1356        assert_eq!(phfv.data(), &[0, 0, 1, 2, 0, 1, 2, 3, 0, 0, 1, 0]);
1357
1358        let mut phf2 = PartitionsHolder::allocated(3, 4, false);
1359        let mut phfv2 = phf2.view();
1360        phfv2.push_partition(&Partition::from("AABC".as_bytes()));
1361        phfv2.push_partition(&Partition::from("ABCD".as_bytes()));
1362        phfv2.push_partition(&Partition::from("AABA".as_bytes()));
1363        assert_eq!(phfv2.n_partitions(), 3);
1364        assert_eq!(phfv.data, phfv2.data);
1365        unsafe {
1366            *phfv2.get_unchecked_mut((0, 1)) = 9;
1367        }
1368        phfv2[(1, 3)] = 9;
1369        assert_eq!(phfv2.data(), &[0, 9, 1, 2, 0, 1, 2, 9, 0, 0, 1, 0]);
1370        assert_eq!(phfv2.get(0), Partition::from("ABCD".as_bytes()));
1371        assert_eq!(phfv2.get(2), Partition::from("AABA".as_bytes()));
1372
1373        let mut pht = PartitionsHolder::allocated(3, 4, true);
1374        assert!(pht.by_row());
1375        assert_eq!(pht.n_partitions(), 3);
1376        let mut phtv = pht.view();
1377        assert_eq!(phtv.data(), &[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);
1378        phtv.push_partition(&Partition::from("AABC".as_bytes()));
1379        phtv.push_partition(&Partition::from("ABCD".as_bytes()));
1380        phtv.push_partition(&Partition::from("AABA".as_bytes()));
1381        assert_eq!(phtv.n_partitions(), 3);
1382        assert_eq!(phtv.data(), &[0, 0, 0, 0, 1, 0, 1, 2, 1, 2, 3, 0]);
1383        unsafe {
1384            *phtv.get_unchecked_mut((0, 1)) = 9;
1385        }
1386        phtv[(1, 3)] = 9;
1387        assert_eq!(phtv.data(), &[0, 0, 0, 9, 1, 0, 1, 2, 1, 2, 9, 0]);
1388        assert_eq!(phtv.get(0), Partition::from("ABCD".as_bytes()));
1389        assert_eq!(phtv.get(2), Partition::from("AABA".as_bytes()));
1390    }
1391}
1392
1393#[doc(hidden)]
1394#[no_mangle]
1395pub unsafe extern "C" fn dahl_partition__enumerated(
1396    n_partitions: i32,
1397    n_items: i32,
1398    partitions_ptr: *mut i32,
1399) {
1400    let n_partitions = usize::try_from(n_partitions).unwrap();
1401    let n_items = usize::try_from(n_items).unwrap();
1402    let mut phv = PartitionsHolderBorrower::from_ptr(partitions_ptr, n_partitions, n_items, true);
1403    for partition in Partition::iter(n_items) {
1404        phv.push_slice(&partition[..]);
1405    }
1406}
1407
1408/// A data structure representation a permutation of integers.
1409///
1410#[derive(Debug)]
1411pub struct Permutation(Vec<usize>);
1412
1413impl Permutation {
1414    pub fn from_slice(x: &[usize]) -> Option<Self> {
1415        let mut y = Vec::from(x);
1416        y.sort_unstable();
1417        for (i, j) in y.into_iter().enumerate() {
1418            if i != j {
1419                return None;
1420            }
1421        }
1422        Some(Self(Vec::from(x)))
1423    }
1424
1425    pub fn from_vector(x: Vec<usize>) -> Option<Self> {
1426        let mut y = x.clone();
1427        y.sort_unstable();
1428        for (i, j) in y.into_iter().enumerate() {
1429            if i != j {
1430                return None;
1431            }
1432        }
1433        Some(Self(x))
1434    }
1435
1436    pub fn natural(n_items: usize) -> Self {
1437        Self((0..n_items).collect())
1438    }
1439
1440    pub fn random<T: Rng>(n_items: usize, rng: &mut T) -> Self {
1441        let mut perm = Self::natural(n_items);
1442        perm.shuffle(rng);
1443        perm
1444    }
1445
1446    pub fn shuffle<T: Rng>(&mut self, rng: &mut T) {
1447        self.0.shuffle(rng)
1448    }
1449
1450    pub fn len(&self) -> usize {
1451        self.0.len()
1452    }
1453
1454    pub fn is_empty(&self) -> bool {
1455        self.0.is_empty()
1456    }
1457
1458    pub fn slice_until(&self, end: usize) -> &[usize] {
1459        &self.0[..end]
1460    }
1461}
1462
1463impl std::ops::Index<usize> for Permutation {
1464    type Output = usize;
1465    fn index(&self, i: usize) -> &Self::Output {
1466        &self.0[i]
1467    }
1468}
1469
1470/// A data structure representing a square matrix.
1471///
1472pub struct SquareMatrix {
1473    data: Vec<f64>,
1474    n_items: usize,
1475}
1476
1477impl SquareMatrix {
1478    pub fn zeros(n_items: usize) -> Self {
1479        Self {
1480            data: vec![0.0; n_items * n_items],
1481            n_items,
1482        }
1483    }
1484
1485    pub fn ones(n_items: usize) -> Self {
1486        Self {
1487            data: vec![1.0; n_items * n_items],
1488            n_items,
1489        }
1490    }
1491
1492    pub fn identity(n_items: usize) -> Self {
1493        let ni1 = n_items + 1;
1494        let n2 = n_items * n_items;
1495        let mut data = vec![0.0; n2];
1496        let mut i = 0;
1497        while i < n2 {
1498            data[i] = 1.0;
1499            i += ni1
1500        }
1501        Self { data, n_items }
1502    }
1503
1504    pub fn data(&self) -> &[f64] {
1505        &self.data[..]
1506    }
1507
1508    pub fn data_mut(&mut self) -> &mut [f64] {
1509        &mut self.data[..]
1510    }
1511
1512    pub fn view(&mut self) -> SquareMatrixBorrower {
1513        SquareMatrixBorrower::from_slice(&mut self.data[..], self.n_items)
1514    }
1515
1516    pub fn n_items(&self) -> usize {
1517        self.n_items
1518    }
1519}
1520
1521pub struct SquareMatrixBorrower<'a> {
1522    data: &'a mut [f64],
1523    n_items: usize,
1524}
1525
1526impl std::ops::Index<(usize, usize)> for SquareMatrixBorrower<'_> {
1527    type Output = f64;
1528    fn index(&self, (i, j): (usize, usize)) -> &Self::Output {
1529        &self.data[self.n_items * j + i]
1530    }
1531}
1532
1533impl std::ops::IndexMut<(usize, usize)> for SquareMatrixBorrower<'_> {
1534    fn index_mut(&mut self, (i, j): (usize, usize)) -> &mut Self::Output {
1535        &mut self.data[self.n_items * j + i]
1536    }
1537}
1538
1539impl<'a> SquareMatrixBorrower<'a> {
1540    pub fn from_slice(data: &'a mut [f64], n_items: usize) -> Self {
1541        assert_eq!(data.len(), n_items * n_items);
1542        Self { data, n_items }
1543    }
1544
1545    /// # Safety
1546    ///
1547    /// Added for FFI.
1548    pub unsafe fn from_ptr(data: *mut f64, n_items: usize) -> Self {
1549        let data = slice::from_raw_parts_mut(data, n_items * n_items);
1550        Self { data, n_items }
1551    }
1552
1553    pub fn n_items(&self) -> usize {
1554        self.n_items
1555    }
1556
1557    /// # Safety
1558    ///
1559    /// You're on your own with the indices.
1560    pub unsafe fn get_unchecked(&self, (i, j): (usize, usize)) -> &f64 {
1561        self.data.get_unchecked(self.n_items * j + i)
1562    }
1563
1564    /// # Safety
1565    ///
1566    /// You're on your own with the indices.
1567    pub unsafe fn get_unchecked_mut(&mut self, (i, j): (usize, usize)) -> &mut f64 {
1568        self.data.get_unchecked_mut(self.n_items * j + i)
1569    }
1570
1571    pub fn data(&self) -> &[f64] {
1572        self.data
1573    }
1574
1575    pub fn data_mut(&mut self) -> &mut [f64] {
1576        self.data
1577    }
1578
1579    pub fn sum_of_triangle(&self) -> f64 {
1580        let mut sum = 0.0;
1581        for i in 0..self.n_items {
1582            for j in 0..i {
1583                sum += unsafe { *self.get_unchecked((i, j)) };
1584            }
1585        }
1586        sum
1587    }
1588
1589    pub fn sum_of_row_subset(&self, row: usize, columns: &[usize]) -> f64 {
1590        let mut sum = 0.0;
1591        for j in columns {
1592            sum += unsafe { *self.get_unchecked((row, *j)) };
1593        }
1594        sum
1595    }
1596}