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

https://leetcode.com/problems/count-the-number-of-arrays-with-k-matching-adjacent-elements/description/?envType=daily-question&envId=2025-06-17

leetcode - Count the Number of Arrays with K Matching Adjacent Elemetns

문제 유형: 조합론, 수학

문제 난이도: Hard

 

문제

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.

 

n, m, k를 입력받는다. 좋은 배열은 다음의 조건을 따른다:

n의 길이에, 1 ~ m의 수의 범위의 요소를 가지고, k개의 요소가 자신의 앞 인덱스의 수와 같아야 한다.

 

구성 가능한 좋은 배열의 개수를 구하시오.

값이 커질 수 있으니, 1e9 + 7로 나눈 나머지를 구하시오.

 

풀이

첫번째 수는 m개의 경우의 수가 있다.

나머지 n-1개 중에서, k개는 자신의 앞과 같아야 한다. 즉, n-1개중 k개를 뽑아야 한다.

그리고 n-k-1개의 수들은 m-1개의 나머지 숫자 중 각각 고유하게 고를 수 있다.

 

즉 정답은 

이다.

 

모듈러 역원 및 페르마 소정리 및 분할정복 거듭제곱을 이용하여 효율적으로 풀 수 있다.

 

코드

C++

#pragma GCC optimize("O3", "unroll-loops");
static const int __ = [](){
    ios::sync_with_stdio(0);
    cin.tie(0);
    return 0;
}();

using ll = long long;
const int MOD = 1e9 + 7;

ll fact[100001], invfact[100001];
bool built = false;

ll bpow(ll base, ll exp) {
    ll res = 1;
    while(exp) {
        if(exp&1) res = (res * base) % MOD;
        exp >>= 1;
        base = (base * base) % MOD;
    }
    return res;
}

void build() {
    if (built) return;
    built = true;
    fact[0] = 1;
    for (int i = 1; i <= 100000; i++) {
        fact[i] = fact[i-1] * i % MOD;
    }
    invfact[100000] = bpow(fact[100000], MOD-2);
    for (int i = 99999; i >= 0; i--) {
        invfact[i] = invfact[i+1] * (i+1) % MOD;
    }
}

ll nCr(ll n, ll r) {
    return ((fact[n] * invfact[n-r]) % MOD * invfact[r]) % MOD;
}

class Solution {
public:
    int countGoodArrays(int n, int m, int k) {
        if(k > n-1) return 0;
        if(n == 1) return k==0?m:0;
        build();
        ll chooseK = nCr(n-1, k);
        ll chooseRest = bpow(m-1, n-k-1);
        return (((m * chooseK) % MOD) * chooseRest) % MOD;
    }
};

 

Go

const MOD int64 = 1_000_000_007
const MAX int = 100000
var fact [MAX+1]int64
var invFact [MAX+1]int64
var built bool = false

func bpow(base, exp int64) int64 {
    res := int64(1)
    for exp > 0 {
        if exp%2 == 1 {
            res = (res * base) % MOD
        }
        base = (base * base) % MOD
        exp >>= 1
    }
    return res
}

func initFact() {
    if built {
        return;
    }
    built = true
    fact[0] = 1
    for i := 1; i <= MAX; i++ {
        fact[i] = (fact[i-1] * int64(i)) % MOD
    } 
    invFact[MAX] = bpow(fact[MAX], MOD-2);
    for i := MAX-1; i >= 0; i-- {
        invFact[i] = invFact[i+1] * int64(i+1) % MOD
    }
}

func nCr(n, r int64) int64 {
    return fact[n] * invFact[n-r] % MOD * invFact[r] % MOD
}

func countGoodArrays(n int, m int, k int) int {
    if k > n-1 { return 0 }
    if n == 1 {
        if k == 0 { return m } else { return 0 }
    }

    initFact()
    chooseK := nCr(int64(n-1), int64(k))
    chooseRest := bpow(int64(m-1), int64(n-k-1))
    return int(int64(m) * chooseK % MOD * chooseRest % MOD)
}

 

Python

MOD = 10**9+7
MAX = 100001

fact = [1] * MAX
invFact = [1] * MAX
built = False

def build():
    global built
    if built:
        return
    built = True
    for i in range(1, MAX):
        fact[i] = fact[i-1] * i % MOD
    invFact[-1] = pow(fact[-1], MOD-2, MOD)
    for i in reversed(range(1, MAX)):
        invFact[i-1] = invFact[i] * (i) % MOD

def nCr(n: int, r: int) -> int:
    return fact[n] * invFact[n-r] % MOD * invFact[r] % MOD

class Solution:
    def countGoodArrays(self, n: int, m: int, k: int) -> int:
        if k > n-1: return 0
        if n == 1: return m if k == 0 else 0
        build()
        return m * nCr(n-1, k) % MOD * pow(m-1, n-k-1, MOD) % MOD