Count the number of arrays with k matching adjacent elements

Math Combinatorics

58.4% acceptance

Feb 25, 2026

429

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 109 + 7.

Solution

Rust

Time O(2^n)

Space O(n)

LeetCode

solution.rs

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 pairs = n - 1;
    if k > pairs {
      return 0;
    }
    fn pow_mod(mut base: i64, mut exp: usize, modulus: i64) -> i64 {
      let mut result = 1i64;
      base %= modulus;
      while exp > 0 {
        if exp & 1 == 1 {
          result = result * base % modulus;
        }
        base = base * base % modulus;
        exp >>= 1;
      }
      result
    }
    let max_n = pairs + 1;
    let mut fact = vec![1i64; max_n + 1];
    for i in 1..=max_n {
      fact[i] = fact[i - 1] * i as i64 % MOD;
    }
    let mut inv_fact = vec![1i64; max_n + 1];
    inv_fact[max_n] = pow_mod(fact[max_n], (MOD - 2) as usize, MOD);
    for i in (0..max_n).rev() {
      inv_fact[i] = inv_fact[i + 1] * (i + 1) as i64 % MOD;
    }
    let comb = fact[pairs] * inv_fact[k] % MOD * inv_fact[pairs - k] % MOD;
    let diff_pow = pow_mod(m - 1, pairs - k, MOD);
    (m % MOD * comb % MOD * diff_pow % MOD) as i32
  }
}