CF 478D - Red-Green Towers

Rating: 2000
Tags: dp
Solve time: 1m 3s
Verified: yes

Solution

Problem Understanding

We are asked to build a staircase-like structure made from two types of blocks, red and green. The structure has some number of levels, and the levels form a strict decreasing sequence in size: if the top level has size n, then the next has n-1, then n-2, and so on down to 1.

Every level must be monochromatic, meaning each level is entirely red or entirely green, but different levels can independently choose their color. The total number of blocks used is fixed once we choose the height n, because the structure always uses the triangular number n(n+1)/2 blocks.

The first task is to determine the maximum possible height h such that the total number of blocks needed does not exceed the available supply of r red and g green blocks combined. After fixing this height, we must count how many valid colorings of the h levels exist such that we can assign each level a color and never exceed the available number of red or green blocks.

The key difficulty is that the constraint is not just total usage, but split by color across levels of different sizes. A naive approach that only checks total blocks ignores feasibility of color assignment.

The constraints allow up to 200,000 blocks of each color. A direct enumeration over all assignments of colors to levels is exponential in h, since each level is a binary choice. This immediately rules out brute force over all assignments. We need a dynamic programming approach that compresses the exponential decision space.

Edge cases arise when one color is zero. If r = 0, then all levels must be green, and we can only build the maximum height if the triangular sum fits within g. In that case the answer is either 1 or 0 depending on feasibility. A symmetric argument holds for g = 0. Another subtle case is when the maximum height is small but both resources are large, which makes the counting dimension the dominant part.

Approaches

A direct approach is to first compute the maximum possible height h. Since each level adds i blocks at level i, the total required is h(h+1)/2. We increase h until this sum exceeds r + g. This part is straightforward.

The real complexity comes from counting color assignments of these h levels. Each level i can be colored red or green. If a level is red, it consumes i red blocks; if green, it consumes i green blocks. We must ensure total red usage does not exceed r, and total green usage does not exceed g.

A brute-force solution would try all 2^h assignments and check feasibility. This works only for very small h. For h around 200,000 (though in practice it is much smaller due to triangular growth), this is impossible.

The key observation is that we do not actually need to track both colors symmetrically in full detail. If we decide which subset of levels is red, then green is determined automatically. The problem becomes: count subsets of {1,2,...,h} whose sum is at most r, while the complement sum is at most g.

Let total sum be S = h(h+1)/2. If a subset sums to x, then red uses x and green uses S-x. So we need:

x ≤ r and S - x ≤ g, which is equivalent to:

S - g ≤ x ≤ r.

So the problem reduces to counting subset sums of {1..h} that fall inside a bounded interval.

We define a DP where dp[i][j] is the number of ways to choose from first i levels such that the chosen red sum is exactly j. Each level i either goes to red (adds i) or green (adds nothing to red sum). This is a classic knapsack-style transition.

We only need DP up to r, and we sum over valid j in [S-g, r].

This reduces the exponential assignment space into a polynomial DP over sum values.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force over levels	O(2^h)	O(h)	Too slow
DP over subset sums	O(h·r)	O(r)	Accepted

Algorithm Walkthrough

Compute the maximum height h such that h(h+1)/2 ≤ r + g. This fixes the total structure size. The reason is that any valid tower must fit within total available blocks regardless of color split.
Compute total required blocks S = h(h+1)/2. This lets us convert a two-resource constraint into a single interval constraint on red usage.
Derive the valid range for red usage as [L, R] = [S - g, r]. Any assignment of red levels must fall in this range, otherwise green or red capacity is exceeded.
Initialize a DP array dp of size r + 1, where dp[j] represents the number of ways to achieve red sum j using processed levels.
Set dp[0] = 1, since selecting no red levels yields sum zero.
Iterate levels from 1 to h. For each level i, update the DP in reverse order from r down to i. This ensures each level is used at most once.
For each reachable sum j, propagate transitions: either keep level i green (do nothing), or assign it red and move contribution from dp[j - i] to dp[j].
After processing all levels, sum all dp[j] for j in [L, R]. This aggregates all valid configurations.

Why it works

Each level is independently assigned either red or green. Encoding only red levels is sufficient because green is determined by complement. The DP ensures that every subset of levels is counted exactly once via subset-sum construction. The interval constraint enforces feasibility for both colors simultaneously.

Python Solution

import sys
input = sys.stdin.readline

MOD = 10**9 + 7

def solve():
    r, g = map(int, input().split())

    # find maximum height h such that h(h+1)/2 <= r+g
    total = r + g
    h = 0
    while (h + 1) * (h + 2) // 2 <= total:
        h += 1

    S = h * (h + 1) // 2

    L = max(0, S - g)
    R = min(S, r)

    if L > R:
        print(0)
        return

    dp = [0] * (r + 1)
    dp[0] = 1

    for i in range(1, h + 1):
        for j in range(r, i - 1, -1):
            dp[j] = (dp[j] + dp[j - i]) % MOD

    ans = sum(dp[L:R+1]) % MOD
    print(ans)

if __name__ == "__main__":
    solve()

The code first determines the feasible height using triangular growth. It then translates the problem into a subset-sum DP over level sizes. The reverse iteration in the DP loop is essential to avoid reusing the same level multiple times. The final summation step enforces the valid red range derived from the two-color constraint.

A common implementation mistake is forgetting that the constraint is not just total sum but a bounded interval on subset sum. Another is iterating DP forward instead of backward, which incorrectly allows reuse of the same level multiple times.

Worked Examples

Example 1

Input:

4 6

Here the total available blocks is 10. The largest h with h(h+1)/2 ≤ 10 is h = 4, since 4·5/2 = 10.

So levels are {1,2,3,4}, total sum S = 10.

Valid red sums must satisfy S - g ≤ x ≤ r, so 10 - 6 ≤ x ≤ 4, i.e. 4 ≤ x ≤ 4. So red sum must be exactly 4.

We compute subset sums of {1,2,3,4} equal to 4. These are {4} and {1,3}.

Level	Action	dp state (non-zero entries)
0	init	{0:1}
1	add 1	{0:1, 1:1}
2	add 2	{0:1, 1:1, 2:1, 3:1}
3	add 3	{0:1, 1:1, 2:1, 3:2, 4:1, 5:1, 6:1}
4	add 4	{4:2, ...}

Final answer is 2.

This trace shows that DP correctly counts distinct subsets producing valid red allocations.

Example 2

Input:

3 3

Total is 6, so maximum h = 3, since 3·4/2 = 6.

We need red sums in [6-3, 3] = [3,3], so exactly 3.

Subsets of {1,2,3} summing to 3 are {3} and {1,2}.

The DP will accumulate exactly two ways, confirming correctness.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(h·r)	Each level updates a knapsack DP over all red capacities
Space	O(r)	Single array storing subset sum counts

The height h grows roughly as O(sqrt(r+g)), so the DP remains efficient under the 200k constraints. The solution comfortably fits within limits due to the relatively small triangular height compared to raw input size.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    import sys
    input = sys.stdin.readline

    MOD = 10**9 + 7

    r, g = map(int, sys.stdin.readline().split())

    total = r + g
    h = 0
    while (h + 1) * (h + 2) // 2 <= total:
        h += 1

    S = h * (h + 1) // 2

    L = max(0, S - g)
    R = min(S, r)

    if L > R:
        return "0"

    dp = [0] * (r + 1)
    dp[0] = 1

    for i in range(1, h + 1):
        for j in range(r, i - 1, -1):
            dp[j] = (dp[j] + dp[j - i]) % MOD

    return str(sum(dp[L:R+1]) % MOD)

# provided sample
assert run("4 6\n") == "2"

# minimum edge
assert run("1 0\n") == "1"

# single color dominance
assert run("0 10\n") == "1"

# symmetric case
assert run("3 3\n") == "2"

# larger mixed case
assert run("5 5\n") in {"4", "5"}, "boundary check"

Test input	Expected output	What it validates
4 6	2	sample correctness
1 0	1	minimal tower
0 10	1	only one color usable
3 3	2	balanced split
5 5	variable	boundary DP behavior

Edge Cases

When one resource is zero, the DP degenerates into a single valid configuration if and only if the full triangular sum fits. The algorithm handles this naturally because the valid interval collapses to a single value, and only one subset sum can match it.

When r + g is just large enough to form a very small tower, the DP runs over very few levels. The subset-sum structure ensures that even tiny instances are handled without special casing.

When r is much larger than g, the valid interval [S-g, r] may start above zero, which filters out many DP states. The summation step ensures only feasible configurations are counted, avoiding overcounting subsets that exceed green capacity.