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 1: Combinatorics + Fast Power

For an array of length \(n\), there are \(n - 1\) pairs of adjacent elements. We need to select \(k\) of these \(n - 1\) adjacent pairs such that the two elements in each of these \(k\) pairs are equal, and the remaining \(n - 1 - k\) adjacent pairs have different elements.

This is equivalent to splitting the array \(n - 1 - k\) times, resulting in \(n - k\) segments, where all elements in each segment are equal. The number of ways to split is \(C_{n - 1}^{n - 1 - k} = C_{n - 1}^{k}\).

For the first segment, we can choose any element from \([1, m]\). For the remaining \(n - k - 1\) segments, we just need to ensure that the element in each segment is different from the previous segment, so each of these segments has \(m - 1\) choices. In total, there are \(m \times (m - 1)^{n - k - 1}\) ways to choose.

Combining the two parts above, we get the answer:

\[ C_{n - 1}^{k} \times m \times (m - 1)^{n - k - 1} \bmod (10^9 + 7) \]

In the code implementation, we can precompute factorials and inverses, and use fast power to calculate combinations.

Ignoring the preprocessing time and space, the time complexity is \(O(\log (n - k))\), and the space complexity is \(O(1)\).

Python3JavaC++GoTypeScriptRust

mx = 10**5 + 10
mod = 10**9 + 7
f = [1] + [0] * mx
g = [1] + [0] * mx

for i in range(1, mx):
    f[i] = f[i - 1] * i % mod
    g[i] = pow(f[i], mod - 2, mod)


def comb(m: int, n: int) -> int:
    return f[m] * g[n] * g[m - n] % mod


class Solution:
    def countGoodArrays(self, n: int, m: int, k: int) -> int:
        return comb(n - 1, k) * m * pow(m - 1, n - k - 1, mod) % mod

class Solution {
    private static final int N = (int) 1e5 + 10;
    private static final int MOD = (int) 1e9 + 7;
    private static final long[] f = new long[N];
    private static final long[] g = new long[N];

    static {
        f[0] = 1;
        g[0] = 1;
        for (int i = 1; i < N; ++i) {
            f[i] = f[i - 1] * i % MOD;
            g[i] = qpow(f[i], MOD - 2);
        }
    }

    public static long qpow(long a, int k) {
        long res = 1;
        while (k != 0) {
            if ((k & 1) == 1) {
                res = res * a % MOD;
            }
            k >>= 1;
            a = a * a % MOD;
        }
        return res;
    }

    public static long comb(int m, int n) {
        return (int) f[m] * g[n] % MOD * g[m - n] % MOD;
    }

    public int countGoodArrays(int n, int m, int k) {
        return (int) (comb(n - 1, k) * m % MOD * qpow(m - 1, n - k - 1) % MOD);
    }
}

const int MX = 1e5 + 10;
const int MOD = 1e9 + 7;
long long f[MX];
long long g[MX];

long long qpow(long a, int k) {
    long res = 1;
    while (k != 0) {
        if ((k & 1) == 1) {
            res = res * a % MOD;
        }
        k >>= 1;
        a = a * a % MOD;
    }
    return res;
}

int init = []() {
    f[0] = g[0] = 1;
    for (int i = 1; i < MX; ++i) {
        f[i] = f[i - 1] * i % MOD;
        g[i] = qpow(f[i], MOD - 2);
    }
    return 0;
}();

long long comb(int m, int n) {
    return f[m] * g[n] % MOD * g[m - n] % MOD;
}

class Solution {
public:
    int countGoodArrays(int n, int m, int k) {
        return comb(n - 1, k) * m % MOD * qpow(m - 1, n - k - 1) % MOD;
    }
};

const MX = 1e5 + 10
const MOD = 1e9 + 7

var f [MX]int64
var g [MX]int64

func qpow(a int64, k int) int64 {
    res := int64(1)
    for k != 0 {
        if k&1 == 1 {
            res = res * a % MOD
        }
        a = a * a % MOD
        k >>= 1
    }
    return res
}

func init() {
    f[0], g[0] = 1, 1
    for i := 1; i < MX; i++ {
        f[i] = f[i-1] * int64(i) % MOD
        g[i] = qpow(f[i], MOD-2)
    }
}

func comb(m, n int) int64 {
    return f[m] * g[n] % MOD * g[m-n] % MOD
}

func countGoodArrays(n int, m int, k int) int {
    ans := comb(n-1, k) * int64(m) % MOD * qpow(int64(m-1), n-k-1) % MOD
    return int(ans)
}

const MX = 1e5 + 10;
const MOD = BigInt(1e9 + 7);

const f: bigint[] = Array(MX).fill(1n);
const g: bigint[] = Array(MX).fill(1n);

function qpow(a: bigint, k: number): bigint {
    let res = 1n;
    while (k !== 0) {
        if ((k & 1) === 1) {
            res = (res * a) % MOD;
        }
        a = (a * a) % MOD;
        k >>= 1;
    }
    return res;
}

(function init() {
    for (let i = 1; i < MX; ++i) {
        f[i] = (f[i - 1] * BigInt(i)) % MOD;
        g[i] = qpow(f[i], Number(MOD - 2n));
    }
})();

function comb(m: number, n: number): bigint {
    return (((f[m] * g[n]) % MOD) * g[m - n]) % MOD;
}

export function countGoodArrays(n: number, m: number, k: number): number {
    const ans = (((comb(n - 1, k) * BigInt(m)) % MOD) * qpow(BigInt(m - 1), n - k - 1)) % MOD;
    return Number(ans);
}

impl Solution {
    pub fn count_good_arrays(n: i32, m: i32, k: i32) -> i32 {
        const N: usize = 1e5 as usize + 10;
        const MOD: i64 = 1_000_000_007;
        use std::sync::OnceLock;

        static F: OnceLock<Vec<i64>> = OnceLock::new();
        static G: OnceLock<Vec<i64>> = OnceLock::new();

        fn qpow(mut a: i64, mut k: i64, m: i64) -> i64 {
            let mut res = 1;
            while k != 0 {
                if k & 1 == 1 {
                    res = res * a % m;
                }
                a = a * a % m;
                k >>= 1;
            }
            res
        }

        fn init() -> (&'static Vec<i64>, &'static Vec<i64>) {
            F.get_or_init(|| {
                let mut f = vec![1i64; N];
                for i in 1..N {
                    f[i] = f[i - 1] * i as i64 % MOD;
                }
                f
            });

            G.get_or_init(|| {
                let f = F.get().unwrap();
                let mut g = vec![1i64; N];
                for i in 1..N {
                    g[i] = qpow(f[i], MOD - 2, MOD);
                }
                g
            });

            (F.get().unwrap(), G.get().unwrap())
        }

        fn comb(f: &[i64], g: &[i64], m: usize, n: usize) -> i64 {
            f[m] * g[n] % MOD * g[m - n] % MOD
        }

        let (f, g) = init();
        let n = n as usize;
        let m = m as i64;
        let k = k as usize;

        let c = comb(f, g, n - 1, k);
        let pow = qpow(m - 1, (n - 1 - k) as i64, MOD);
        (c * m % MOD * pow % MOD) as i32
    }
}

Comments