Problem

You are given three integers n, m, k. A good array arr of size n is defined as follows:

  • Each element in arr is in the inclusive range [1, m].
  • Exactly k indices i (where 1 <= i < n) satisfy the condition arr[i - 1] == arr[i].

Return the number of good arrays that can be formed.

Since the answer may be very large, return it modulo 10^9 + 7.

Examples

Example 1

1
2
3
4
5

Input: n = 3, m = 2, k = 1
Output: 4
Explanation:
  * There are 4 good arrays. They are `[1, 1, 2]`, `[1, 2, 2]`, `[2, 1, 1]` and `[2, 2, 1]`.
  * Hence, the answer is 4.

Example 2

1
2
3
4
5

Input: n = 4, m = 2, k = 2
Output: 6
Explanation:
  * The good arrays are `[1, 1, 1, 2]`, `[1, 1, 2, 2]`, `[1, 2, 2, 2]`, `[2, 1, 1, 1]`, `[2, 2, 1, 1]` and `[2, 2, 2, 1]`.
  * Hence, the answer is 6.

Example 3

1
2
3
4

Input: n = 5, m = 2, k = 0
Output: 2
Explanation:
  * The good arrays are `[1, 2, 1, 2, 1]` and `[2, 1, 2, 1, 2]`. Hence, the answer is 2.

Constraints

  • 1 <= n <= 10^5
  • 1 <= m <= 10^5
  • 0 <= k <= n - 1

Solution

Method 1 – Dynamic Programming with Combinatorics

Intuition

We want to count arrays of length n with elements from 1 to m such that exactly k adjacent pairs are equal. The key is to realize that the array can be split into k+1 blocks of consecutive equal numbers, and each block must have a different value from its neighbor. The number of ways to partition the array and assign values can be computed using combinatorics and fast exponentiation.

Approach

  1. The array can be split into k+1 blocks (since each matching adjacent pair reduces the number of blocks by 1).
  2. The number of ways to choose the positions of the k matching pairs among n-1 possible adjacent pairs is C(n-1, k).
  3. The first block can be any of m values. Each subsequent block must be different from the previous, so there are (m-1) choices for each.
  4. Multiply the number of ways to partition (C(n-1, k)) by the number of ways to assign values (m * (m-1)^k).
  5. Use modular arithmetic for large numbers.
C++
 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

class Solution {
  const int MOD = 1e9 + 7;
  vector<long long> fac, inv;
  long long modpow(long long a, long long b) {
    long long res = 1;
    while (b) {
      if (b & 1) res = res * a % MOD;
      a = a * a % MOD;
      b >>= 1;
    }
    return res;
  }
  void init(int n) {
    fac.resize(n+1, 1);
    inv.resize(n+1, 1);
    for (int i = 1; i <= n; ++i) fac[i] = fac[i-1] * i % MOD;
    inv[n] = modpow(fac[n], MOD-2);
    for (int i = n-1; i >= 0; --i) inv[i] = inv[i+1] * (i+1) % MOD;
  }
  long long comb(int n, int k) {
    if (k < 0 || k > n) return 0;
    return fac[n] * inv[k] % MOD * inv[n-k] % MOD;
  }
public:
  int countGoodArrays(int n, int m, int k) {
    init(n);
    long long ans = comb(n-1, k) * m % MOD * modpow(m-1, n-1-k) % MOD;
    return ans;
  }
};
Go
 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

func modpow(a, b, mod int) int {
  res := 1
  for b > 0 {
    if b&1 == 1 {
      res = res * a % mod
    }
    a = a * a % mod
    b >>= 1
  }
  return res
}
func countGoodArrays(n, m, k int) int {
  mod := int(1e9 + 7)
  fac := make([]int, n)
  inv := make([]int, n)
  fac[0] = 1
  for i := 1; i < n; i++ {
    fac[i] = fac[i-1] * i % mod
  }
  inv[n-1] = modpow(fac[n-1], mod-2, mod)
  for i := n-2; i >= 0; i-- {
    inv[i] = inv[i+1] * (i+1) % mod
  }
  comb := func(a, b int) int {
    if b < 0 || b > a {
      return 0
    }
    return fac[a] * inv[b] % mod * inv[a-b] % mod
  }
  ans := comb(n-1, k) * m % mod * modpow(m-1, n-1-k, mod) % mod
  return ans
}
Java
 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

class Solution {
  static final int MOD = (int)1e9 + 7;
  long[] fac, inv;
  long modpow(long a, long b) {
    long res = 1;
    while (b > 0) {
      if ((b & 1) == 1) res = res * a % MOD;
      a = a * a % MOD;
      b >>= 1;
    }
    return res;
  }
  void init(int n) {
    fac = new long[n+1];
    inv = new long[n+1];
    fac[0] = 1;
    for (int i = 1; i <= n; ++i) fac[i] = fac[i-1] * i % MOD;
    inv[n] = modpow(fac[n], MOD-2);
    for (int i = n-1; i >= 0; --i) inv[i] = inv[i+1] * (i+1) % MOD;
  }
  long comb(int n, int k) {
    if (k < 0 || k > n) return 0;
    return fac[n] * inv[k] % MOD * inv[n-k] % MOD;
  }
  public int countGoodArrays(int n, int m, int k) {
    init(n);
    long ans = comb(n-1, k) * m % MOD * modpow(m-1, n-1-k) % MOD;
    return (int)ans;
  }
}
Kotlin
 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

class Solution {
  private val MOD = 1_000_000_007L
  private lateinit var fac: LongArray
  private lateinit var inv: LongArray

  private fun modpow(a: Long, b: Long): Long {
    var res = 1L
    var aa = a
    var bb = b
    while (bb > 0) {
      if (bb and 1L == 1L) res = res * aa % MOD
      aa = aa * aa % MOD
      bb = bb shr 1
    }
    return res
  }

  private fun init(n: Int) {
    fac = LongArray(n + 1) { 1L }
    inv = LongArray(n + 1) { 1L }
    for (i in 1..n) fac[i] = fac[i - 1] * i % MOD
    inv[n] = modpow(fac[n], MOD - 2)
    for (i in n - 1 downTo 0) inv[i] = inv[i + 1] * (i + 1) % MOD
  }

  private fun comb(n: Int, k: Int): Long {
    if (k < 0 || k > n) return 0L
    return fac[n] * inv[k] % MOD * inv[n - k] % MOD
  }

  fun countGoodArrays(n: Int, m: Int, k: Int): Int {
    init(n)
    val ans = comb(n - 1, k) * m % MOD * modpow((m - 1).toLong(), (n - 1 - k).toLong()) % MOD
    return ans.toInt()
  }
}
Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

class Solution:
  MOD = 10**9 + 7

  def countGoodArrays(self, n: int, m: int, k: int) -> int:
    fac = [1] * (n)
    inv = [1] * (n)
    for i in range(1, n):
      fac[i] = fac[i-1] * i % self.MOD
    inv[-1] = pow(fac[-1], self.MOD-2, self.MOD)
    for i in range(n-2, -1, -1):
      inv[i] = inv[i+1] * (i+1) % self.MOD
    def comb(a: int, b: int) -> int:
      if b < 0 or b > a:
        return 0
      return fac[a] * inv[b] % self.MOD * inv[a-b] % self.MOD
    ans = comb(n-1, k) * m % self.MOD * pow(m-1, n-1-k, self.MOD) % self.MOD
    return ans
Rust
 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

impl Solution {
  pub fn count_good_arrays(n: i32, m: i32, k: i32) -> i32 {
    const MOD: i64 = 1_000_000_007;
    let n = n as usize;
    let m = m as i64;
    let k = k as usize;
    let mut fac = vec![1i64; n];
    let mut inv = vec![1i64; n];
    for i in 1..n {
      fac[i] = fac[i-1] * (i as i64) % MOD;
    }
    inv[n-1] = pow(fac[n-1], MOD-2, MOD);
    for i in (0..n-1).rev() {
      inv[i] = inv[i+1] * ((i+1) as i64) % MOD;
    }
    fn pow(mut a: i64, mut b: i64, m: i64) -> i64 {
      let mut res = 1;
      while b > 0 {
        if b & 1 == 1 {
          res = res * a % m;
        }
        a = a * a % m;
        b >>= 1;
      }
      res
    }
    fn comb(fac: &Vec<i64>, inv: &Vec<i64>, n: usize, k: usize, m: i64) -> i64 {
      if k > n { return 0; }
      fac[n] * inv[k] % m * inv[n-k] % m
    }
    let ans = comb(&fac, &inv, n-1, k, MOD) * m % MOD * pow(m-1, (n-1-k) as i64, MOD) % MOD;
    ans as i32
  }
}

Complexity

  • ⏰ Time complexity: O(n) (for precomputing factorials and inverses)
  • 🧺 Space complexity: O(n) (for factorial and inverse arrays)