rust-partitions

partitions.rs

//-----------------------------------------
//-----------------------------------------
//-----------------------------------------
// These functions are for unordered sequences (meaning that
//   e.g.  2 + 3 and 3 + 2 are considered to be the same solution,
//         and thus only one is returned)
//-----------------------------------------

// Unique ways (up to order) to add any amount of numbers x (min <= x <= max)
// to produce a specific sum.
//
// When min=1 and max=sum, the number of solutions that this should produce
// is given by https://oeis.org/A000041 :
//
//          n:  0  1  2  3  4  5   6   7   8   9  10  11  12   13 ...
// f(n).len():  1  1  2  3  5  7  11  15  22  30  42  56  77  101 ...
fn unordered_partitions(sum: i64, min: i64, max: i64) -> Vec<Vec<i64>> {
	(0..sum+1)
		.map(|k| unordered_k_partitions(sum, min, max, k))
		.fold(vec![], |mut b, mut x| { b.append(&mut x); b})
}

// Unique ways (up to order) to add 'nterms' numbers in (min...max)
// to a given sum.
fn unordered_k_partitions(sum: i64, min: i64, max: i64, nterms: i64) -> Vec<Vec<i64>> {
	if nterms == 0 {
		if sum == 0 { return vec![vec![]]; } // success base case
		else { return vec![]; } // failure base case
	} else {
		let max = ::std::cmp::min(sum - min * (nterms - 1), max);
		(min..max+1).flat_map(|x|
			unordered_k_partitions(sum-x, min, x, nterms-1)
				.into_iter().map(move |mut vec| { vec.push(x); vec })
		).collect()
	}
}

//-----------------------------------------
//-----------------------------------------
//-----------------------------------------
// These functions are for ordered sequences (meaning that
//   e.g.  2 + 3 and 3 + 2 are considered distinct solutions)
//-----------------------------------------

// Unique ways to add any amount of numbers x (min <= x <= max)
// to produce a specific sum.
//
// The number of solutions that this *should* produce
//  is given by this sequence:  http://oeis.org/A011782
//
//          n:  0  1  2  3  4   5   6   7    8    9   10  ...
// f(n).len():  1  1  2  4  8  16  32  64  128  256  512  ...
fn ordered_partitions(sum: i64) -> Vec<Vec<i64>> {
	(0..sum+1)
		.map(|k| ordered_k_partitions(sum, k))
		.fold(vec![], |mut b, mut x| { b.append(&mut x); b})
}


// Unique ways to add 'nterms' numbers in (min...max)
// to a given sum.
fn ordered_k_partitions(sum: i64, k: i64) -> Vec<Vec<i64>> {
	assert!(k >= 0);
	if k == 0 {
		if sum == 0 { return vec![vec![]]; } // success base case
		else { return vec![]; } // failure base case
	}
	(1..sum+1).flat_map(|x|
		ordered_k_partitions(sum-x, k-1)
			.into_iter().map(move |mut vec| { vec.push(x); vec })
	).collect()
}