random_partition/
lib.rs

1//! Generate approximately uniformly distributed random integer partitions.
2//! So given natural numbers n, k find a sequence of natural (nonzero) p₁, ..., pₖ such
3//! that n = ∑ᵢ₌₁ᵏ pᵢ.
4use ndarray::{s, Array2, Array3};
5use num::integer::div_floor;
6use rug::Integer;
7
8pub fn number_of_partitions_into_parts(total: usize, number_of_parts: usize) -> Integer {
9    match (number_of_parts, total) {
10        (0, 0) => 1.into(),
11        (0, _) => 0.into(),
12        (_, _) => {
13            // Should realistically be able to do this with just two rows - but eh
14            // TODO: rewrite this to use only two rows acting as a ring buffer
15            let mut counts: Array2<Integer> = Array2::zeros((number_of_parts + 1, total + 1));
16            counts[[0, 0]] = 1.into();
17            for n in 1..=total {
18                counts[[1, n]] = 1.into();
19            }
20            for k in 2..=number_of_parts {
21                counts[[k, k]] = 1.into();
22                for n in k + 1..=total {
23                    counts[[k, n]] = (&counts[[k, n - k]] + &counts[[k - 1, n - 1]]).into();
24                }
25            }
26            std::mem::take(&mut counts[[number_of_parts, total]])
27        }
28    }
29}
30
31pub fn number_of_partitions_into_parts_with_max(
32    total: usize,
33    number_of_parts: usize,
34    max_part: usize,
35) -> Array3<Integer> {
36    let mut counts: Array3<Integer> = Array3::zeros((total + 1, number_of_parts + 1, max_part + 1));
37    for m in 0..=max_part {
38        counts[[0, 0, m]] = 1.into();
39        // s
40        for k in 1..=number_of_parts {
41            for n in k..=std::cmp::min(m * k, total) {
42                for i in 0..=div_floor(n, m) {
43                    let x = counts[[n - i * m, k - i, m - 1]].clone();
44                    counts[[n, k, m]] += x;
45                }
46            }
47        }
48    }
49    counts
50}
51
52fn part_lower_bound(total: usize, number_of_parts: usize) -> usize {
53    let m1 = div_floor(total, number_of_parts);
54    if total % number_of_parts != 0 {
55        m1 + 1
56    } else {
57        m1
58    }
59}
60
61fn random_partition_buf_from_table<'a, R: rug::rand::MutRandState>(
62    rng: &mut R,
63    total: usize,
64    number_of_parts: usize,
65    buf: &'a mut [usize],
66    constrained_part_table: &Array3<Integer>,
67) -> &'a mut [usize] {
68    let mut remaining_total = total;
69    let number_of_possible_partitions = number_of_partitions_into_parts(total, number_of_parts);
70    let mut which: Integer = number_of_possible_partitions.random_below(rng) + 1; // rng.gen_range(1.into()..=number_of_possible_partitions); // TODO: random int from 1..=number_of_possible_partitions
71    let mut ub = total - number_of_parts + 1;
72    let mut lb = part_lower_bound(remaining_total, number_of_parts);
73    for (i, remaining_parts) in (1..=number_of_parts).rev().enumerate() {
74        let mut count = &constrained_part_table[[remaining_total, remaining_parts, lb]];
75        let mut part_size = lb;
76        'l: for k in lb..=ub {
77            count = &constrained_part_table[[remaining_total, remaining_parts, k]];
78            part_size = k;
79            if count >= &which {
80                count = &constrained_part_table[[remaining_total, remaining_parts, k - 1]];
81                break 'l;
82            }
83        }
84        buf[i] = part_size;
85        remaining_total -= part_size;
86        if remaining_total == 0 {
87            break;
88        }
89        which -= count;
90        lb = part_lower_bound(remaining_total, remaining_parts - 1);
91        ub = part_size;
92    }
93    &mut buf[0..number_of_parts]
94}
95
96/// Generates a random partition of `total` into `number_of_parts` writing the result into the provided buffer.
97///
98/// # Returns
99/// A slice into the input buffer now containing a random partition with elements ordered in
100/// descending order.
101///
102/// # Arguments
103/// * `rng` - the random number generator used
104/// * `total` - the number being partitioned
105/// * `number_of_parts` - how many pieces the partition should have
106/// * `buf` - a buffer with at least enough space for `number_of_parts` integers that the
107///     returned values are written into.
108pub fn random_partition_buf<'a, R: rug::rand::MutRandState>(
109    rng: &mut R,
110    total: usize,
111    number_of_parts: usize,
112    buf: &'a mut [usize],
113) -> &'a mut [usize] {
114    let ub = total - number_of_parts + 1;
115    random_partition_buf_from_table(
116        rng,
117        total,
118        number_of_parts,
119        buf,
120        &number_of_partitions_into_parts_with_max(total, number_of_parts, ub),
121    )
122}
123
124/// Generates a random partition of `total` into `number_of_parts`.
125///
126/// # Returns
127/// A vector containing a random partition with elements ordered in descending order.
128///
129/// # Arguments
130/// * `rng` - the random number generator used
131/// * `total` - the number being partitioned
132/// * `number_of_parts` - how many pieces the partition should have
133pub fn random_partition<R: rug::rand::MutRandState>(
134    rng: &mut R,
135    total: usize,
136    number_of_parts: usize,
137) -> Vec<usize> {
138    let mut partition = vec![0; number_of_parts];
139    random_partition_buf(rng, total, number_of_parts, &mut partition);
140    partition
141}
142
143/// Generates multiple random partitions of `total` into `number_of_parts`.
144///
145/// # Returns
146/// A 2D array where each row represents a partition with elements ordered in descending order.
147///
148/// # Arguments
149/// * `rng` - the random number generator used
150/// * `total` - the number being partitioned
151/// * `number_of_parts` - how many pieces the partition should have
152/// * `number_of_partitions` - the total number of partitions to generate
153pub fn random_partitions<R: rug::rand::MutRandState>(
154    rng: &mut R,
155    total: usize,
156    number_of_parts: usize,
157    number_of_partitions: usize,
158) -> Array2<usize> {
159    // This array has to be in row major order - this is the documented standard behaviour of ndarray
160    let mut partitions = Array2::zeros((number_of_partitions, number_of_parts));
161    let ub = total - number_of_parts + 1;
162    let table = number_of_partitions_into_parts_with_max(total, number_of_parts, ub);
163    for i in 0..number_of_partitions {
164        random_partition_buf_from_table(
165            rng,
166            total,
167            number_of_parts,
168            partitions.slice_mut(s![i, ..]).as_slice_mut().unwrap(),
169            &table,
170        );
171    }
172    partitions
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    #[test]
180    fn partition_with_max() {
181        let sol = number_of_partitions_into_parts_with_max(200, 61, 53);
182        assert_eq!(sol[[200, 61, 53]], Integer::from(13253620047_u64));
183        assert_eq!(sol[[200, 55, 53]], Integer::from(23898693166_u64));
184        assert_eq!(sol[[103, 55, 24]], Integer::from(139935_u64));
185        assert_eq!(sol[[103, 5, 24]], Integer::from(119_u64));
186    }
187
188    #[test]
189    fn random() {
190        let mut rng = rug::rand::RandState::new();
191        dbg!(random_partition(&mut rng, 500, 23));
192        // panic!()
193    }
194
195    #[test]
196    fn random_parts() {
197        let mut rng = rug::rand::RandState::new();
198        dbg!(random_partitions(&mut rng, 500, 23, 10));
199        // panic!()
200    }
201}