3405. Count the Number of Arrays with K Matching Adjacent Elements

Description

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.

 

Example 1:

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:

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:

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 <= 105
• 1 <= m <= 105
• 0 <= k <= n - 1

 

Solutions

Solution: Combinatorics

• Time complexity: O(nlogn)
• Space complexity: O(n)

 

JavaScript

/**
 * @param {number} n
 * @param {number} m
 * @param {number} k
 * @return {number}
 */
const MODULO = BigInt(10 ** 9 + 7);
const MAX = 10 ** 5;
const fact = new Array(MAX).fill(1n);
const invFact = new Array(MAX).fill(1n);

function modPow(base, exponent) {
  let result = 1n;

  base = BigInt(base);
  exponent = BigInt(exponent);

  while (exponent > 0n) {
    if (exponent % 2n) {
      result = (result * base) % MODULO;
    }

    base = (base * base) % MODULO;
    exponent /= 2n;
  }

  return result;
}

for (let index = 1; index < MAX; index++) {
  fact[index] = (fact[index - 1] * BigInt(index)) % MODULO;
}

invFact[MAX - 1] = modPow(fact[MAX - 1], MODULO - 2n);

for (let index = MAX - 2; index >= 0; index--) {
  invFact[index] = (invFact[index + 1] * BigInt(index + 1)) % MODULO;
}

const countGoodArrays = function (n, m, k) {
  const total = (BigInt(m) * modPow(m - 1, n - k - 1)) % MODULO;

  const nCk = (n, k) => {
    if (k < 0 || k > n) return 0n;

    return (((fact[n] * invFact[k]) % MODULO) * invFact[n - k]) % MODULO;
  };

  return Number((total * nCk(n - 1, k)) % MODULO);
};