LeetCode 1692 - Count Ways to Distribute Candies

The problem asks us to compute the number of ways to distribute n distinct candies into k bags such that every bag has a

LeetCode Problem 1692

Difficulty: 🔴 Hard
Topics: Dynamic Programming

Solution

Problem Understanding

The problem asks us to compute the number of ways to distribute n distinct candies into k bags such that every bag has at least one candy. The candies are unique and labeled from 1 to n, and the order of the candies within a bag or the order of the bags themselves does not matter. Two distributions are considered different if the set of candies in at least one bag differs between the two distributions.

The input consists of two integers, n and k. The output is a single integer representing the number of valid distributions modulo $10^9 + 7$. The constraints $1 \le k \le n \le 1000$ indicate that the solution must be efficient, likely requiring dynamic programming or combinatorial mathematics rather than brute-force enumeration, because enumerating all partitions grows exponentially.

Key points to note include: all candies must be distributed, every bag must contain at least one candy, and the order does not matter. Edge cases include when n == k, where each bag must get exactly one candy, and k == 1, where all candies go into a single bag.

Approaches

A brute-force approach would attempt to generate all partitions of n elements into k non-empty subsets. This is theoretically correct but computationally infeasible because the number of ways grows rapidly (related to Stirling numbers of the second kind). Enumerating all possible distributions would be exponential in time complexity and thus not feasible for n up to 1000.

The key insight for an optimal solution is recognizing that this problem is equivalent to computing the Stirling numbers of the second kind, which counts the number of ways to partition n distinct items into k non-empty subsets. The recursive formula for Stirling numbers is:

$$S(n, k) = k \cdot S(n-1, k) + S(n-1, k-1)$$

The reasoning is: when adding the n-th candy, it can either go into one of the existing k bags (hence k * S(n-1, k)) or start a new bag (hence S(n-1, k-1)). Using dynamic programming, we can compute all S(n, k) values efficiently and return the result modulo $10^9 + 7$.

Approach Time Complexity Space Complexity Notes
Brute Force O(k^n) O(n) Enumerate all distributions explicitly. Exponential and infeasible.
Optimal (DP/Stirling Numbers) O(n * k) O(n * k) Use recursive DP with memoization or iterative table to compute Stirling numbers modulo 10^9+7.

Algorithm Walkthrough

  1. Initialize a 2D DP table dp of size (n+1) x (k+1) where dp[i][j] represents the number of ways to distribute i candies into j bags.
  2. Set the base case dp[0][0] = 1, meaning there is one way to distribute zero candies into zero bags.
  3. For each number of candies i from 1 to n:
  • For each number of bags j from 1 to k:

  • Compute dp[i][j] using the recursive relation dp[i][j] = (j * dp[i-1][j] + dp[i-1][j-1]) % MOD.

  • j * dp[i-1][j] accounts for placing the new candy into one of the existing j bags.

  • dp[i-1][j-1] accounts for placing the new candy into a new bag.

  1. Return dp[n][k] as the final answer modulo $10^9 + 7$.

Why it works: The recurrence ensures that all partitions are counted exactly once, and the modulo operation prevents integer overflow. The base case and recursive construction cover all possible distributions while satisfying the constraint that each bag contains at least one candy.

Python Solution

class Solution:
    def waysToDistribute(self, n: int, k: int) -> int:
        MOD = 10**9 + 7
        dp = [[0] * (k+1) for _ in range(n+1)]
        dp[0][0] = 1

        for i in range(1, n+1):
            for j in range(1, k+1):
                dp[i][j] = (j * dp[i-1][j] + dp[i-1][j-1]) % MOD

        return dp[n][k]

The implementation uses a 2D array dp to store intermediate results. We iterate through each possible number of candies and bags, applying the Stirling number recurrence to fill the table. Using modulo ensures correctness under large numbers.

Go Solution

func waysToDistribute(n int, k int) int {
    const MOD = 1000000007
    dp := make([][]int, n+1)
    for i := range dp {
        dp[i] = make([]int, k+1)
    }
    dp[0][0] = 1

    for i := 1; i <= n; i++ {
        for j := 1; j <= k; j++ {
            dp[i][j] = (j*dp[i-1][j] + dp[i-1][j-1]) % MOD
        }
    }

    return dp[n][k]
}

The Go implementation mirrors the Python logic. We initialize a slice of slices to represent the DP table, iterate through candies and bags, and apply the recurrence. Integer multiplication and addition are carefully handled modulo $10^9 + 7$ to prevent overflow.

Worked Examples

Example 1: n = 3, k = 2

i j dp[i][j]
0 0 1
1 1 1
2 1 1 * 1 + 0 = 1
2 2 1 * 0 + 1 = 1
3 1 1 * 1 + 0 = 1
3 2 2 * 1 + 1 = 3

Output: 3

Example 2: n = 4, k = 2

Step-by-step filling yields dp[4][2] = 7.

Example 3: n = 20, k = 5

DP table computation modulo $10^9 + 7$ gives dp[20][5] = 206085257.

Complexity Analysis

Measure Complexity Explanation
Time O(n * k) We iterate over all i from 1 to n and j from 1 to k to fill the DP table.
Space O(n * k) The DP table stores n * k entries for intermediate results.

The DP approach is efficient for the constraints n, k <= 1000. Each entry is computed in constant time.

Test Cases

s = Solution()

# Provided examples
assert s.waysToDistribute(3, 2) == 3  # Example 1
assert s.waysToDistribute(4, 2) == 7  # Example 2
assert s.waysToDistribute(20, 5) == 206085257  # Example 3

# Edge cases
assert s.waysToDistribute(1, 1) == 1  # smallest possible n and k
assert s.waysToDistribute(5, 1) == 1  # all candies in one bag
assert s.waysToDistribute(5, 5) == 1  # each bag gets exactly one candy
assert s.waysToDistribute(10, 2) == 511  # moderate values
Test Why
(3, 2) Simple case with small numbers, verifies basic correctness
(4, 2) Verifies correctness with multiple distribution possibilities
(20, 5) Tests modulo handling with large numbers
(1, 1) Minimal input, edge case
(5, 1) All candies in one bag, simplest distribution
(5, 5) Each candy in its own bag, edge case where n == k
(10, 2) Moderate size to check DP correctness

Edge Cases

Edge Case 1: n == k

When the number of candies equals the number of bags, each bag must receive exactly one candy. The DP recurrence correctly handles this because dp[i][i] builds from dp[i-1][i-1], ensuring all candies are placed into separate bags.

Edge Case 2: k == 1

All candies go into a single bag. The recurrence works because dp[i][1] multiplies 1 * dp[i-1][1], effectively counting only one distribution.

Edge Case 3: Large n, k near n

This tests efficiency and modulo handling. Since the DP uses a table of size n * k and computes modulo $10^9 + 7$,