Atcoder — M — Candies

5 min readMay 18, 2024

https://atcoder.jp/contests/dp/tasks/dp_m

Problem

We have given n children starting from index 0 to n-1. Array a[i] represents the capacity for the i children to eat the chocolate. The goal is to find out how many ways to distribute k chocolates to children to make sure no chocolates are left.

Example

a = {1,2,3}
k = 4

Let us try to manually solve it.

 0  1  2
{x}{y}{z} where x+y+z= k
{0}{0}{4} -- 2nd violated
{0}{1}{3}
{0}{2}{2}
{0}{3}{1} -- 1st violated
{0}{4}{0} -- 1st violated
{1}{0}{3}
{1}{1}{2}
{1}{2}{1}
{1}{3}{0} -- 1st violated
{2}{0}{2} -- rest at 0th will violate 
Total = 5

Bruteforce approach (TLE)

We can apply the manual approach we tried above. I.e. for each, f(i,k) we have 0..j..k choices. Once ith child chooses the chocolate as per constraints the remaining k-j chocolate will be left over for the rest of the children to try. We can write the recursive function as below.

f(i,k) = sum(f(i+1, k-j)) where j is from 0 to k and i represent i....n-1

We can write the following code to implement the above recursion function.

int recursion(int i, int k) {
    // base case1: zero chocolates can be distributed in 1 ways to all children
    if(k == 0) {
        return 1;
    }
    // base case2: no children then there would be 0 ways to distribute k chocolates
    if(i == n) {
        return 0;    
    }  
    int ans = 0;
    // ith children will try all possible chocolates
    for (int j=0; j <=k; j++) {
        // check to make sure child eat chocolates less or equal to their capacity
        int capacity = a[i];
        if(j <= capacity) {
            ans += recursion(i+1, k - j);
        }
    }
    return ans;
}

The time complexity of this program would be O(k^n) . Which is exponential and therefore impossible to run this code on large input.

Recursion with memoization (TLE)

Since f(i,k) repeat itself so we can apply memoization as below.

vector<vector<int>> dp (n+1, vector<int>(k+1, -1));

int recursion(int i, int k) {
    // base case1: zero chocolates can be distributed in 1 ways to all children
    if(k == 0) {
        return 1;
    }
    // base case2: no children then there would be 0 ways to distribute k chocolates
    if(i == n) {
        return 0;    
    }  
    // if ans exist in dp then return early
    if(dp[i][k] != -1) {
        return dp[i][k];
    }
    int ans = 0;
    // ith children will try all possible chocolates
    for (int j=0; j <=k; j++) {
        // check to make sure child eat chocolates less or equal to their capacity
        int capacity = a[i];
        if(j <= capacity) {
            ans += recursion(i+1, k - j);
        }
    }
    dp[i][k] = ans;
    return ans;
}

The time complexity of this code is O(n*k*k) . This code is optimized but will still throw a TLE error.

Bottom-up DP (TLE)

The same recursion formula can be written to start from the bottom up i.e. solve problems of smaller size starting from 0 and slowly solve the bigger size.

f(i,k) = sum(f(i-1, k-j)) where j is from 0 to k and i represent 0...i

We can code it as below.

    vector<vector<int>> dp (n+1, vector<int>(k+1, -1));
    // Base case: 0 chocolate can be distributed to any number of children in 1 way
    for (int i=0; i < n+1; i++) {
        dp[i][0] = 1;
    }
    // Base case: 0 children with > 0 chocolates can be distributed in 0 ways
    for (int j=1; j < k + 1; j++) {
        dp[0][j] = 0;
    }    

    // fill rest 
    for (int i=1; i < n + 1; i++) {
        for (int j=1; j < k + 1; j++) {
            int ans = 0;
            // Try all possible chocolates
            for (int x=0; x < j + 1; x++) {
                // check to make sure child eat x chocolates less or equal to their capacity
                int capacity = a[i-1];
                if(x <= capacity) {
                    // rest of chocolates i.e. j - x will be available for rest
                    ans += dp[i-1][j - x];
                }
            }
            dp[i][j] = ans;
        }
    }
    cout<<dp[n][k]<<endl;

The time complexity of this code is O(n*k*k) . This code is optimized but will still throw a TLE error.

Optimized DP with prefix sum (works)

How can we improve bottom DP further?

`ans += dp[i-1][j — x]` represent sum of distributing x chocolate to ith child and remaining j-x chocolates to i-1 children, where x range from 0 from j
ans is technically the sum of previous rows i-1 from 0 to j if capacity >= j .
If capacity < j it means ith child can not take all j chocolate. It can max take capacity only. It also means for i-1 row, we can allocate a minimum j-capacity (when ith will take capacity)and max can be j (when ith will take 0 ). It means we need to find the sum of rows for i-1 between [j-capacity,j].
This can be calculated using the prefix sum.
prefix sum for array nums of size n is defined as below

prefix[0] = 0;
for (int i=1; i < n; i++) {
  prefix[i] = prefix[i-1] + nums[i];
}

To find out the sum between sub-array (i,j) can be queried over prefix as below.

sum(i,j) = nums[i]+.....+nums[j]
sum(i,j) = prefix[j] - prefix[i-1]

Applying the same to our problem here, we can optimize calculating sum using prefix sum as below.
if capacity >= j then need to find out sum from 0 to j using ans=prefix[i-1][j]
else need to find out sum from [j-capacity, j] using ans=prefix[j]-prefix[j-capacity-1]

Note: Here prefix is 2d to represent the prefix sum for each ith row. It can also be converted into1d to optimize the space utilization. But for now, I will use the 2d to showcase the solution.

Following would be updated code.

    vector<vector<int>> dp(n+1, vector<int>(k+1,0)); 
    vector<vector<int>> prefix(n+1, vector<int>(k+1,0)); 
    // Base case: 0 chocolate can be distributed to any number of children in 1 way
    for (int i=0; i < n+1; i++) {
        dp[i][0] = 1;
    }
    // Base case: 0 children with > 0 chocolates can be distributed in 0 ways
    for (int j=1; j < k + 1; j++) {
        dp[0][j] = 0;
    }
    
    // fill rest 
    for (int i=1; i < n + 1; i++) {
        for (int j=1; j < k + 1; j++) {
            int capacity = a[i-1] ;
            if(capacity >= j) {
              dp[i][j] = prefix[i-1][j];  
            } else {
                dp[i][j] = (prefix[i-1][j] -  prefix[i-1][j - capacity - 1] + PRIME)%PRIME;  
            }
            // update prefix at runtime
            prefix[i][j] = (prefix[i][j-1] + dp[i][j])%PRIME;
        }
    }
    cout<<dp[n][k]<<endl;

The time complexity of this code is O(n*k) .

Note: I have also added PRIME in the above code to match the problem condition to avoid overflow.

Happy coding :-)