1420. Build Array Where You Can Find The Maximum Exactly K Comparisons

Approach 1: w/o Prefix Sum

• Time: $O(nm^k)$
• Space: $O(nmk)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 class Solution { public: int numOfArrays(int n, int m, int k) { constexpr int kMod = 1'000'000'007; // dp[i][j][k] := # of ways to build an array of length i, where j is the // max used num and k is the search_cost vector>> dp( n + 1, vector>(m + 1, vector(k + 1))); for (int j = 1; j <= m; ++j) dp[1][j][1] = 1; for (int i = 2; i <= n; ++i) // for each length for (int j = 1; j <= m; ++j) // for each max value for (int cost = 1; cost <= k; ++cost) { // for each cost // 1. appending any of [1, j] in i-th position // doesn't change the max and cost dp[i][j][cost] = static_cast(j) * dp[i - 1][j][cost] % kMod; // 2. appending j in i-th position // make j the new max and cost 1 for (int prevMax = 1; prevMax < j; ++prevMax) { dp[i][j][cost] += dp[i - 1][prevMax][cost - 1]; dp[i][j][cost] %= kMod; } } int ans = 0; for (int j = 1; j <= m; ++j) { ans += dp[n][j][k]; ans %= kMod; } return ans; } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 class Solution { public int numOfArrays(int n, int m, int k) { final int kMod = 1_000_000_007; // dp[i][j][k] := # of ways to build an array of length i, where j is the // max num and k is the search_cost int[][][] dp = new int[n + 1][m + 1][k + 1]; for (int j = 1; j <= m; ++j) dp[1][j][1] = 1; for (int i = 2; i <= n; ++i) // for each length for (int j = 1; j <= m; ++j) // for each max value for (int cost = 1; cost <= k; ++cost) { // for each cost // 1. appending any of [1, j] in i-th position // doesn't change the max and cost dp[i][j][cost] = (int) ((long) j * dp[i - 1][j][cost] % kMod); // 2. appending j in i-th position // make j the new max and cost 1 for (int prevMax = 1; prevMax < j; ++prevMax) { dp[i][j][cost] += dp[i - 1][prevMax][cost - 1]; dp[i][j][cost] %= kMod; } } int ans = 0; for (int j = 1; j <= m; ++j) { ans += dp[n][j][k]; ans %= kMod; } return ans; } } 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 class Solution: def numOfArrays(self, n: int, m: int, k: int) -> int: kMod = 1_000_000_007 # dp[i][j][k] := # of ways to build an array of length i, where j is the # max used num and k is the search_cost dp = [[[0] * (k + 1) for j in range(m + 1)] for _ in range(n + 1)] for j in range(1, m + 1): dp[1][j][1] = 1 for i in range(2, n + 1): # for each length for j in range(1, m + 1): # for each max value for cost in range(1, k + 1): # for each cost # 1. appending any of [1, j] in i-th position # doesn't change the max and cost dp[i][j][cost] = j * dp[i - 1][j][cost] % kMod # 2. appending j in i-th position # make j the new max and cost 1 for prevMax in range(1, j): dp[i][j][cost] += dp[i - 1][prevMax][cost - 1] dp[i][j][cost] %= kMod return sum(dp[n][j][k] for j in range(1, m + 1)) % kMod 

Approach 2: w/ Prefix Sum

• Time: $O(nmk)$
• Space: $O(nmk)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 class Solution { public: int numOfArrays(int n, int m, int k) { constexpr int kMod = 1'000'000'007; // dp[i][j][k] := # of ways to build an array of length i, where j is the // max num and k is the search_cost vector>> dp( n + 1, vector>(m + 1, vector(k + 1))); // prefix[i][j][k] := sum(dp[i][x][k]), where 1 <= x <= j vector>> prefix( n + 1, vector>(m + 1, vector(k + 1))); for (int j = 1; j <= m; ++j) { dp[1][j][1] = 1; prefix[1][j][1] = j; } for (int i = 2; i <= n; ++i) // for each length for (int j = 1; j <= m; ++j) // for each max value for (int cost = 1; cost <= k; ++cost) { // for each cost // 1. appending any of [1, j] in i-th position // doesn't change the max and cost // 2. appending j in i-th position // make j the new max and cost 1 dp[i][j][cost] = (static_cast(j) * dp[i - 1][j][cost] + prefix[i - 1][j - 1][cost - 1]) % kMod; prefix[i][j][cost] = (dp[i][j][cost] + prefix[i][j - 1][cost]) % kMod; } int ans = 0; for (int j = 1; j <= m; ++j) { ans += dp[n][j][k]; ans %= kMod; } return ans; } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 class Solution { public int numOfArrays(int n, int m, int k) { final int kMod = 1_000_000_007; // dp[i][j][k] := # of ways to build an array of length i, where j is the // max num and k is the search_cost int[][][] dp = new int[n + 1][m + 1][k + 1]; // prefix[i][j][k] := sum(dp[i][x][k]), where 1 <= x <= j int[][][] prefix = new int[n + 1][m + 1][k + 1]; for (int j = 1; j <= m; ++j) { dp[1][j][1] = 1; prefix[1][j][1] = j; } for (int i = 2; i <= n; ++i) // for each length for (int j = 1; j <= m; ++j) // for each max value for (int cost = 1; cost <= k; ++cost) { // for each cost // 1. appending any of [1, j] in i-th position // doesn't change the max and cost // 2. appending j in i-th position // make j the new max and cost 1 dp[i][j][cost] = (int) (((long) j * dp[i - 1][j][cost] + prefix[i - 1][j - 1][cost - 1]) % kMod); prefix[i][j][cost] = (dp[i][j][cost] + prefix[i][j - 1][cost]) % kMod; } int ans = 0; for (int j = 1; j <= m; ++j) { ans += dp[n][j][k]; ans %= kMod; } return ans; } } 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 class Solution: def numOfArrays(self, n: int, m: int, k: int) -> int: kMod = 1_000_000_007 # dp[i][j][k] := # of ways to build an array of length i, where j is the # max used num and k is the search_cost dp = [[[0] * (k + 1) for j in range(m + 1)] for _ in range(n + 1)] # prefix[i][j][k] := sum(dp[i][x][k]), where 1 <= x <= j prefix = [[[0] * (k + 1) for j in range(m + 1)] for _ in range(n + 1)] for j in range(1, m + 1): dp[1][j][1] = 1 prefix[1][j][1] = j for i in range(2, n + 1): # for each length for j in range(1, m + 1): # for each max value for cost in range(1, k + 1): # for each cost # 1. appending any of [1, j] in i-th position # doesn't change the max and cost # 2. appending j in i-th position # make j the new max and cost 1 dp[i][j][cost] = (j * dp[i - 1][j][cost] + prefix[i - 1][j - 1][cost - 1]) % kMod prefix[i][j][cost] = (dp[i][j][cost] + prefix[i][j - 1][cost]) % kMod return sum(dp[n][j][k] for j in range(1, m + 1)) % kMod