1692. Count Ways to Distribute Candies

There are n unique candies (labeled 1 through n) and k bags. You are asked to distribute all the candies into the bags such that every bag has at least one candy.

There can be multiple ways to distribute the candies. Two ways are considered different if the candies in one bag in the first way are not all in the same bag in the second way. The order of the bags and the order of the candies within each bag do not matter.

For example, (1), (2,3) and (2), (1,3) are considered different because candies 2 and 3 in the bag (2,3) in the first way are not in the same bag in the second way (they are split between the bags (2) and (1,3)). However, (1), (2,3) and (3,2), (1) are considered the same because the candies in each bag are all in the same bags in both ways.

Given two integers, n and k, return the number of different ways to distribute the candies. As the answer may be too large, return it modulo 109 + 7.

 

Example 1:

Input: n = 3, k = 2
Output: 3
Explanation: You can distribute 3 candies into 2 bags in 3 ways:
(1), (2,3)
(1,2), (3)
(1,3), (2)

Example 2:

Input: n = 4, k = 2
Output: 7
Explanation: You can distribute 4 candies into 2 bags in 7 ways:
(1), (2,3,4)
(1,2), (3,4)
(1,3), (2,4)
(1,4), (2,3)
(1,2,3), (4)
(1,2,4), (3)
(1,3,4), (2)

Example 3:

Input: n = 20, k = 5
Output: 206085257
Explanation: You can distribute 20 candies into 5 bags in 1881780996 ways. 1881780996 modulo 109 + 7 = 206085257.

 

Constraints:

  • 1 <= k <= n <= 1000

Rust Solution

struct Solution;

use std::cmp::Ordering::*;

const MOD: i64 = 1_000_000_007;

impl Solution {
    fn ways_to_distribute(n: i32, k: i32) -> i32 {
        let n = n as usize;
        let k = k as usize;
        let mut memo: Vec<Vec<Option<i64>>> = vec![vec![None; k + 1]; n + 1];
        Self::dp(n, k, &mut memo) as i32
    }

    fn dp(n: usize, k: usize, memo: &mut Vec<Vec<Option<i64>>>) -> i64 {
        match n.cmp(&k) {
            Equal => 1,
            Less => 0,
            Greater => {
                if k == 0 {
                    0
                } else {
                    if let Some(res) = memo[n][k] {
                        res
                    } else {
                        let res = (k as i64 * Self::dp(n - 1, k, memo)
                            + Self::dp(n - 1, k - 1, memo))
                            % MOD;
                        memo[n][k] = Some(res);
                        res
                    }
                }
            }
        }
    }
}

#[test]
fn test() {
    let n = 3;
    let k = 2;
    let res = 3;
    assert_eq!(Solution::ways_to_distribute(n, k), res);
    let n = 4;
    let k = 2;
    let res = 7;
    assert_eq!(Solution::ways_to_distribute(n, k), res);
    let n = 20;
    let k = 5;
    let res = 206085257;
    assert_eq!(Solution::ways_to_distribute(n, k), res);
}

Having problems with this solution? Click here to submit an issue on github.