Skip to content

1692. Count Ways to Distribute Candies 👍

Approach 1: 2D DP

  • Time: $O(mn)$
  • Space: $O(mn)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
class Solution {
 public:
  int waysToDistribute(int n, int k) {
    constexpr int kMod = 1'000'000'007;
    // dp[i][j] := the number of ways to distribute i candies to j bags
    // dp[i][j] = 0, if i < j
    //          = 1, if i == j
    //          = dp[i - 1][j - 1] (the new candy occupies a bag)
    //          + dp[i - 1][j] * j (the new candy is in any of the j bags)
    vector<vector<long>> dp(n + 1, vector<long>(k + 1));

    for (int i = 0; i <= k; ++i)
      dp[i][i] = 1;

    for (int i = 1; i <= n; ++i)
      for (int j = 1; j <= k; ++j)
        dp[i][j] = (dp[i - 1][j - 1] + dp[i - 1][j] * j) % kMod;

    return dp[n][k];
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
  public int waysToDistribute(int n, int k) {
    final int kMod = 1_000_000_007;
    // dp[i][j] := the number of ways to distribute i candies to j bags
    // dp[i][j] = 0, if i < j
    //          = 1, if i == j
    //          = dp[i - 1][j - 1] (the new candy occupies a bag)
    //          + dp[i - 1][j] * j (the new candy is in any of the j bags)
    long[][] dp = new long[n + 1][k + 1];

    for (int i = 0; i <= k; ++i)
      dp[i][i] = 1;

    for (int i = 1; i <= n; ++i)
      for (int j = 1; j <= k; ++j)
        dp[i][j] = (dp[i - 1][j - 1] + dp[i - 1][j] * j) % kMod;

    return (int) dp[n][k];
  }
}

Approach 2: Optimize the dimensions

  • Time: $O(mn)$
  • Space: $O(mn)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Solution {
 public:
  int waysToDistribute(int n, int k) {
    constexpr int kMod = 1'000'000'007;
    vector<vector<long>> dp(k + 1, vector<long>(n + 1));

    for (int i = 0; i <= k; ++i)
      dp[i][i] = 1;

    for (int i = 1; i <= k; ++i)
      for (int j = i + 1; j <= n; ++j)
        dp[i][j] = (dp[i - 1][j - 1] + i * dp[i][j - 1]) % kMod;

    return dp[k][n];
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
class Solution {
  public int waysToDistribute(int n, int k) {
    final int kMod = 1_000_000_007;
    long[][] dp = new long[k + 1][n + 1];

    for (int i = 0; i <= k; ++i)
      dp[i][i] = 1;

    for (int i = 1; i <= k; ++i)
      for (int j = i + 1; j <= n; ++j)
        dp[i][j] = (dp[i - 1][j - 1] + i * dp[i][j - 1]) % kMod;

    return (int) dp[k][n];
  }
}