Program to count the number of ways to distribute n number of candies in k number of bags in Python

Suppose we have n candies and k bags, and we need to distribute the candies such that each bag contains at least one candy. Since every candy is unique, we must count all possible ways to distribute them among the bags.

This is a classic problem that can be solved using Stirling numbers of the second kind, which count the number of ways to partition n distinct objects into k non-empty subsets, multiplied by k! to account for the distinct bags.

Problem Understanding

For n = 3 candies and k = 2 bags, the possible distributions are ?

Bag 1: {1, 2}, Bag 2: {3}
Bag 1: {1, 3}, Bag 2: {2}  
Bag 1: {2, 3}, Bag 2: {1}

So the answer is 3 ways.

Algorithm Explanation

We use dynamic programming where dp[c][b] represents the number of ways to distribute c candies into b bags such that each bag has at least one candy.

The recurrence relation is ?

• dp[c][b] = dp[c-1][b-1] + dp[c-1][b] * b

• First term: place the c-th candy in a new bag

• Second term: place the c-th candy in one of the existing b bags

Implementation

def count_candy_distributions(n, k):
    # If more bags than candies, impossible to fill each bag
    if k > n:
        return 0
    
    # dp[i][j] = ways to distribute i candies in j bags (each bag non-empty)
    dp = [[0] * (k + 1) for _ in range(n + 1)]
    
    # Base case: 0 ways to put candies in 0 bags (except 0 candies in 0 bags = 1 way)
    dp[0][0] = 1
    
    # Fill the dp table
    for candies in range(1, n + 1):
        for bags in range(1, min(candies, k) + 1):
            # Put current candy in a new bag
            dp[candies][bags] += dp[candies - 1][bags - 1]
            # Put current candy in one of existing bags
            dp[candies][bags] += dp[candies - 1][bags] * bags
    
    return dp[n][k]

# Test with the example
n, k = 3, 2
result = count_candy_distributions(n, k)
print(f"Ways to distribute {n} candies in {k} bags: {result}")

Ways to distribute 3 candies in 2 bags: 3

Step-by-Step Trace

For n=3, k=2, let's trace the DP table ?

def count_candy_distributions_with_trace(n, k):
    dp = [[0] * (k + 1) for _ in range(n + 1)]
    dp[0][0] = 1
    
    print("DP Table Construction:")
    print("dp[candies][bags] = ways to distribute")
    
    for candies in range(1, n + 1):
        for bags in range(1, min(candies, k) + 1):
            dp[candies][bags] = dp[candies - 1][bags - 1] + dp[candies - 1][bags] * bags
            print(f"dp[{candies}][{bags}] = dp[{candies-1}][{bags-1}] + dp[{candies-1}][{bags}] * {bags}")
            print(f"          = {dp[candies - 1][bags - 1]} + {dp[candies - 1][bags]} * {bags} = {dp[candies][bags]}")
    
    return dp[n][k]

result = count_candy_distributions_with_trace(3, 2)
print(f"\nFinal answer: {result}")

DP Table Construction:
dp[candies][bags] = ways to distribute
dp[1][1] = dp[0][0] + dp[0][1] * 1
          = 1 + 0 * 1 = 1
dp[2][1] = dp[1][0] + dp[1][1] * 1
          = 0 + 1 * 1 = 1
dp[2][2] = dp[1][1] + dp[1][2] * 2
          = 1 + 0 * 2 = 1
dp[3][1] = dp[2][0] + dp[2][1] * 1
          = 0 + 1 * 1 = 1
dp[3][2] = dp[2][1] + dp[2][2] * 2
          = 1 + 1 * 2 = 3

Final answer: 3

Alternative Approach Using Stirling Numbers

This problem is equivalent to computing Stirling numbers of the second kind S(n,k) multiplied by k! if bags are distinguishable ?

import math

def stirling_second_kind(n, k):
    """Compute Stirling number of second kind S(n,k)"""
    if k > n or k == 0:
        return 0
    if k == 1 or k == n:
        return 1
    
    # Using the recurrence: S(n,k) = k*S(n-1,k) + S(n-1,k-1)
    dp = [[0] * (k + 1) for _ in range(n + 1)]
    
    # Base cases
    for i in range(n + 1):
        dp[i][0] = 0 if i > 0 else 1
    for j in range(k + 1):
        dp[0][j] = 0 if j > 0 else 1
    
    for i in range(1, n + 1):
        for j in range(1, min(i, k) + 1):
            dp[i][j] = j * dp[i-1][j] + dp[i-1][j-1]
    
    return dp[n][k]

# For distinguishable bags, multiply by k!
n, k = 3, 2
stirling_num = stirling_second_kind(n, k)
total_ways = stirling_num * math.factorial(k)

print(f"Stirling S({n},{k}) = {stirling_num}")
print(f"Total ways with distinguishable bags = {stirling_num} × {k}! = {total_ways}")

Stirling S(3,2) = 3
Total ways with distinguishable bags = 3 × 2! = 6

Comparison

Conclusion

The problem of distributing n unique candies into k bags (each non-empty) uses dynamic programming with the recurrence dp[c][b] = dp[c-1][b-1] + dp[c-1][b] × b. This approach efficiently computes Stirling numbers of the second kind, which count such partitions.