Build Array Where You Can Find The Maximum Exactly K Comparisons: 3D Dynamic Programming

Given three integers n, m, and k, count the number of arrays of length n where every element is between 1 and m (inclusive), and the "search cost" of the array is exactly k. Return the answer modulo 10^9 + 7.

This is LeetCode 1420: Build Array Where You Can Find The Maximum Exactly K Comparisons.

The problem

The search cost of an array is the number of times a new maximum is encountered when you scan from left to right. More precisely, start with a running maximum of negative infinity. Walk through the array one element at a time. Every time the current element is strictly greater than the running maximum, update the running maximum and increment the cost by 1. The final cost is the search cost.

For example, in the array [2, 1, 3, 1, 2], the running maximum goes from nothing to 2 at index 0 (cost becomes 1), stays at 2 through index 1, then jumps to 3 at index 2 (cost becomes 2). Indices 3 and 4 do not exceed 3, so the final search cost is 2.

Your goal is to count how many such arrays exist for given n, m, and k.

Approach: 3D DP on position, cost, and current maximum

The brute force approach would enumerate all m^n arrays and check each one, but that is far too slow. Instead, you can build the array one position at a time and track three things: how many elements you have placed, how many times the maximum has increased so far, and what the current maximum is.

Define dp[i][j][maxVal] as the number of ways to fill the first i positions such that the search cost is exactly j and the current running maximum is maxVal.

When you place the next element, two things can happen:

1. The new element does not exceed the current max. You can pick any value from 1 to maxVal. The search cost stays the same, and the maximum stays the same. This gives maxVal choices, so dp[i+1][j][maxVal] += dp[i][j][maxVal] * maxVal.

2. The new element exceeds the current max. You pick some value newMax strictly greater than maxVal. The search cost increases by 1, and the maximum becomes newMax. For each newMax from maxVal + 1 to m, you get dp[i+1][j+1][newMax] += dp[i][j][maxVal].

Step-by-step walkthrough

The code

def num_of_arrays(n, m, k):
    MOD = 10**9 + 7
    dp = [[[0] * (m + 1) for _ in range(k + 1)] for _ in range(n + 1)]

    for v in range(1, m + 1):
        dp[1][1][v] = 1

    for i in range(1, n):
        for j in range(1, k + 1):
            for maxVal in range(1, m + 1):
                if dp[i][j][maxVal] == 0:
                    continue
                cnt = dp[i][j][maxVal]
                dp[i + 1][j][maxVal] = (dp[i + 1][j][maxVal] + cnt * maxVal) % MOD
                for newMax in range(maxVal + 1, m + 1):
                    dp[i + 1][j + 1][newMax] = (dp[i + 1][j + 1][newMax] + cnt) % MOD

    result = 0
    for maxVal in range(1, m + 1):
        result = (result + dp[n][k][maxVal]) % MOD
    return result

The outer loop walks through positions 1 to n - 1. For each existing state dp[i][j][maxVal], the code pushes contributions forward to position i + 1. The first branch handles elements that do not create a new maximum (multiply by maxVal because there are that many valid choices). The second branch iterates over each possible new maximum and increments the search cost.

After filling the entire table, the answer is the sum of dp[n][k][maxVal] across all values of maxVal from 1 to m.

Complexity analysis

The base 3D DP solution has a nested loop over newMax inside the main triple loop, giving O(n * k * m^2). You can optimize the inner loop with a prefix sum to eliminate one factor of m, bringing it down to O(n * k * m). The space is O(n * k * m) for the 3D table. You can also reduce the position dimension to two layers since each position only depends on the previous one, cutting space to O(k * m).

The building blocks

3D state tracking

Many counting problems need more than two dimensions. Here, two dimensions (position and cost) are not enough because you also need to know the current maximum to determine which transitions are valid. Adding the third dimension (current max) lets each state fully describe what matters for future choices. This pattern appears whenever the "constraint" on the sequence depends on a running aggregate like a maximum, minimum, or sum.

Forward DP (push-style transitions)

Instead of computing each state by looking backward at its predecessors, this solution pushes each state forward to its successors. The push style is natural here because one source state fans out to many target states: maxVal choices that keep the max the same, and up to m - maxVal choices that raise the max. Push-style DP avoids redundant iteration and makes the transition logic cleaner.

Edge cases

k = 0: A search cost of zero is impossible because the first element always establishes a new maximum. The answer is always 0 when k = 0.

k > m: You can never have more than m new-maximum events because each event raises the maximum by at least 1, and the maximum cannot exceed m. If k > m, the answer is 0.

n = 1: With a single element, the search cost is always 1. If k = 1, the answer is m (each value from 1 to m is a valid single-element array). If k != 1, the answer is 0.

m = 1: Every element must be 1, so the running maximum never increases after the first element. The search cost is always 1. If k = 1, the answer is 1. Otherwise, 0.

From understanding to recall

You now know how to set up the 3D DP table, why three dimensions are necessary, and how the two transition cases (keep the max vs. raise the max) partition all possible next elements. But knowing is not the same as doing. Under interview pressure, you need to remember the state definition, the base case (first element always costs 1), and the two branches of the transition without hesitation.

The 3D counting DP pattern shows up in several other problems, including Number of Music Playlists and Dice Roll Simulation. In each case, you identify what "summary" of the sequence so far is enough to determine future transitions, then build dimensions around that summary. Practicing these problems together helps you see the shared skeleton: define states, set base cases, push or pull transitions, and collect the answer.

The 3D DP with forward transitions is one of the reusable building blocks in the CodeBricks system. You practice each block with spaced repetition until the pattern is automatic. When you can write the state definition, base case, and transitions from memory, you own this entire family of counting problems.

Related posts

• Number of Music Playlists - Another counting DP with constraints on distinct elements
• Dice Roll Simulation - DP with multiple state dimensions tracking sequence constraints
• Count Vowels Permutation - Counting sequences with transition rules