1use 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 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 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; 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
96pub 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
124pub 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
143pub 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 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 }
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 }
201}