Skip to content

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

Approach 1: w/o 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
class Solution {
 public:
  int numOfArrays(int n, int m, int k) {
    constexpr int kMod = 1'000'000'007;
    // dp[i][j][k] := the number of ways to build an array of length i, where j
    // is the maximum number and k is `search_cost`
    vector<vector<vector<int>>> dp(
        n + 1, vector<vector<int>>(m + 1, vector<int>(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 the i-th position doesn't change the
          //    maximum and cost.
          dp[i][j][cost] = static_cast<long>(j) * dp[i - 1][j][cost] % kMod;
          // 2. Appending j in the i-th position makes 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
class Solution {
  public int numOfArrays(int n, int m, int k) {
    final int kMod = 1_000_000_007;
    // dp[i][j][k] := the number of ways to build an array of length i, where j
    // is the maximum number and k is `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 the i-th position doesn't change the
          //    maximum and cost.
          dp[i][j][cost] = (int) ((long) j * dp[i - 1][j][cost] % kMod);
          // 2. Appending j in the i-th position makes 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
class Solution:
  def numOfArrays(self, n: int, m: int, k: int) -> int:
    kMod = 1_000_000_007
    # dp[i][j][k] := the number of ways to build an array of length i, where j
    # is the maximum number and k is `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 the i-th position doesn't change the
          #    maximum and cost.
          dp[i][j][cost] = j * dp[i - 1][j][cost] % kMod
          # 2. Appending j in the i-th position makes 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] := the number of ways to build an array of length i, where j
    // is the maximum number and k is the `search_cost`
    vector<vector<vector<int>>> dp(
        n + 1, vector<vector<int>>(m + 1, vector<int>(k + 1)));
    // prefix[i][j][k] := sum(dp[i][x][k]), where 1 <= x <= j
    vector<vector<vector<int>>> prefix(
        n + 1, vector<vector<int>>(m + 1, vector<int>(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 the i-th position doesn't change the
          //    maximum and cost.
          // 2. Appending j in the i-th position makes j the new max and
          //    cost 1.
          dp[i][j][cost] = (static_cast<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
26
27
28
29
30
31
32
33
class Solution {
  public int numOfArrays(int n, int m, int k) {
    final int kMod = 1_000_000_007;
    // dp[i][j][k] := the number of ways to build an array of length i, where j
    // is the maximum number 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 the i-th position doesn't change the
          //    maximum and cost.
          // 2. Appending j in the i-th position makes 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
class Solution:
  def numOfArrays(self, n: int, m: int, k: int) -> int:
    kMod = 1_000_000_007
    # dp[i][j][k] := the number of ways to build an array of length i, where j
    # is the maximum number 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 the i-th position doesn't change the
          #    maximum and cost.
          # 2. Appending j in the i-th position makes 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