Build Array Where You Can Find The Maximum Exactly K Comparisons — Hard Problem & Solution

Consider this way of finding a maximum: Count the arrays of length n whose values lie in [1, m] and for which this procedure ends with cost exactly k.

Difficulty: Hard
Topics: Dynamic Programming, Prefix Sum
Asked at: Amazon, Google, Microsoft
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

Consider this way of finding a maximum:

maximum = -1; cost = 0
for each value in arr:
    if value > maximum:
        maximum = value
        cost = cost + 1

Count the arrays of length n whose values lie in [1, m] and for which this procedure ends with cost exactly k. Return the count modulo 10⁹ + 7.

Example 1

Input: n = 2, m = 3, k = 1
Output: 6
Explanation: The first element must already be the maximum, so the second is any value at most the first.

Example 2

Input: n = 5, m = 2, k = 3
Output: 0
Explanation: With only two distinct values the cost can never reach 3.

Example 3

Input: n = 9, m = 1, k = 1
Output: 1
Explanation: Only the all-ones array.

Constraints

1 <= n <= 50
1 <= m <= 100
0 <= k <= n

How to solve Build Array Where You Can Find The Maximum Exactly K Comparisons

dp[j][c] counts the arrays built so far whose running maximum is j and whose cost is c. Extending by one element either keeps the maximum — j choices, cost unchanged — or introduces a new maximum v > j, which costs one more.

Approach

Answer 0 immediately when k = 0: a non-empty array always costs at least 1.
Seed dp[j][1] = 1 for every j from 1 to m — the first element is always a new maximum.
For each further length, fold dp[j][c] · j into next[j][c] and dp[j][c] into next[v][c+1] for every v > j.
Sum dp[j][k] over j.

Why it works

Tracking the maximum rather than the elements is what makes the count finite and small: everything below the maximum is interchangeable, so "any of j values" collapses a whole branch into a multiplication. The inner loop over v > j is what a prefix-sum optimisation removes, taking the solution from O(n · m² · k) to O(n · m · k) — at these bounds either is comfortable, but the prefix-sum form is what the problem is really testing.

Complexity

Time — O(n · m² · k), or O(n · m · k) with suffix sums
Space — O(m · k)

Pitfalls

k = 0 is impossible for n >= 1 — the first element always triggers an update.
The multiplier is j, the maximum's value, not the number of elements so far.
Costs above k must be discarded, not clamped.

Reference solution

Python

def numOfArrays(n: int, m: int, k: int) -> int:
    MOD = 10**9 + 7
    if k == 0:
        return 0
    dp = [[0] * (k + 1) for _ in range(m + 1)]
    for j in range(1, m + 1):
        dp[j][1] = 1
    for _ in range(2, n + 1):
        nxt = [[0] * (k + 1) for _ in range(m + 1)]
        for j in range(1, m + 1):
            for c in range(1, k + 1):
                ways = dp[j][c]
                if ways == 0:
                    continue
                nxt[j][c] = (nxt[j][c] + ways * j) % MOD
                if c + 1 <= k:
                    for v in range(j + 1, m + 1):
                        nxt[v][c + 1] = (nxt[v][c + 1] + ways) % MOD
        dp = nxt
    return sum(dp[j][k] for j in range(1, m + 1)) % MOD

JavaScript

var numOfArrays = function(n, m, k) {
    var MOD = 1000000007;
    if (k === 0) return 0;
    var j, c, v;
    var make = function() {
        var t = [];
        for (var a = 0; a <= m; a++) {
            var row = [];
            for (var b = 0; b <= k; b++) row.push(0);
            t.push(row);
        }
        return t;
    };
    var dp = make();
    for (j = 1; j <= m; j++) dp[j][1] = 1;
    for (var len = 2; len <= n; len++) {
        var next = make();
        for (j = 1; j <= m; j++) {
            for (c = 1; c <= k; c++) {
                var ways = dp[j][c];
                if (ways === 0) continue;
                next[j][c] = (next[j][c] + ways * j) % MOD;
                if (c + 1 <= k) {
                    for (v = j + 1; v <= m; v++) next[v][c + 1] = (next[v][c + 1] + ways) % MOD;
                }
            }
        }
        dp = next;
    }
    var answer = 0;
    for (j = 1; j <= m; j++) answer = (answer + dp[j][k]) % MOD;
    return answer;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 195 dynamic programming problems · the whole catalogue