LeetCode Problem 3405

Difficulty: 🔴 Hard
Topics: Math, Combinatorics

Solution

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

Problem Understanding

We are given three integers:

• n, the length of the array.
• m, the number of possible values each element may take, namely integers in the range [1, m].
• k, the exact number of adjacent pairs that must be equal.

An array arr is considered good if exactly k indices i with 1 ≤ i < n satisfy:

arr[i - 1] == arr[i]

There are exactly n - 1 adjacent pairs in an array of length n. For each adjacent pair, one of two things happens:

• The two elements are equal.
• The two elements are different.

The task is to count how many arrays of length n satisfy the requirement that exactly k of these adjacent pairs are equal.

Since the answer can be extremely large, we return it modulo:

10^9 + 7

The constraints are quite large:

• n ≤ 100000
• m ≤ 100000

This immediately rules out any approach that attempts to generate arrays explicitly. Even dynamic programming with states proportional to n × m would be too expensive.

Several edge cases are important:

• k = 0, meaning every adjacent pair must be different.
• k = n - 1, meaning all adjacent pairs must be equal.
• m = 1, meaning only one value exists.
• Very large values of n and m, requiring an efficient combinatorial solution.

Approaches

Brute Force

A direct approach would generate all possible arrays of length n.

There are:

m^n

possible arrays.

For each generated array, we could count how many adjacent pairs are equal and increment the answer if the count equals k.

This approach is obviously correct because it examines every possible array and filters exactly those satisfying the condition.

Unfortunately, it is completely infeasible. Even for moderate values such as:

n = 20
m = 10

the number of arrays is:

10^20

which is astronomically large.

Key Insight

Instead of constructing arrays directly, think about the positions where the value changes.

There are n - 1 adjacent gaps:

arr[0]|arr[1]|arr[2]|...|arr[n-1]

For each gap, either:

• The next value equals the previous one.
• The next value differs from the previous one.

If exactly k adjacent pairs are equal, then exactly:

(n - 1) - k

adjacent pairs are different.

Let:

d = (n - 1) - k

We first choose which of the n - 1 gaps are different:

C(n - 1, d)

ways.

Next:

• Choose the first element: m choices.
• Every chosen "different" gap forces the next value to differ from the current one, giving m - 1 choices.
• Every "equal" gap has exactly 1 choice.

Therefore:

Answer = m × C(n - 1, d) × (m - 1)^d

Since:

d = n - 1 - k

we obtain:

Answer = m × C(n - 1, k) × (m - 1)^(n - 1 - k)

because:

C(n - 1, d) = C(n - 1, k)

This transforms the problem into computing a binomial coefficient and a modular exponentiation.

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(m^n · n)	O(n)	Enumerates every array
Optimal	O(n)	O(n)	Uses combinatorics and modular arithmetic

Algorithm Walkthrough

1. Let MOD = 1_000_000_007.
2. Compute the number of unequal adjacent pairs:

d = n - 1 - k

3. Precompute factorials from 0 to n - 1 modulo MOD.
4. Precompute inverse factorials using Fermat's Little Theorem.
5. Compute:

C(n - 1, k)

using:

factorial[n-1] *
inverse_factorial[k] *
inverse_factorial[n-1-k]

6. Compute:

(m - 1)^d mod MOD

using fast modular exponentiation. 7. Multiply the three independent components:

m
× C(n - 1, k)
× (m - 1)^d

8. Return the result modulo MOD.

Why it works

Every valid array can be uniquely described by two independent choices:

1. Which adjacent positions are equal and which are different.
2. What values are assigned while respecting those choices.

Choosing the locations of the differing adjacent pairs determines the array's structure. Once the first element is chosen, each differing gap contributes exactly m - 1 choices and each equal gap contributes exactly one choice. Since every valid array corresponds to exactly one such construction and every construction yields a valid array, the counting formula is exact.

Python Solution

class Solution:
    def countGoodArrays(self, n: int, m: int, k: int) -> int:
        MOD = 1_000_000_007

        fact = [1] * n
        for i in range(1, n):
            fact[i] = fact[i - 1] * i % MOD

        inv_fact = [1] * n
        inv_fact[n - 1] = pow(fact[n - 1], MOD - 2, MOD)

        for i in range(n - 1, 0, -1):
            inv_fact[i - 1] = inv_fact[i] * i % MOD

        comb = (
            fact[n - 1]
            * inv_fact[k]
            % MOD
            * inv_fact[n - 1 - k]
            % MOD
        )

        diff_count = n - 1 - k

        return (
            m
            * comb
            % MOD
            * pow(m - 1, diff_count, MOD)
            % MOD
        )

The implementation directly follows the combinatorial formula.

The factorial array allows efficient computation of binomial coefficients modulo MOD. The inverse factorial array is built using Fermat's Little Theorem because MOD is prime.

After computing:

C(n - 1, k)

the code calculates the number of differing adjacent pairs:

n - 1 - k

and raises m - 1 to that power using Python's built in modular exponentiation.

Finally, all factors are multiplied modulo MOD.

Go Solution

func countGoodArrays(n int, m int, k int) int {
	const MOD int64 = 1000000007

	modPow := func(base, exp int64) int64 {
		result := int64(1)
		base %= MOD

		for exp > 0 {
			if exp&1 == 1 {
				result = result * base % MOD
			}
			base = base * base % MOD
			exp >>= 1
		}

		return result
	}

	fact := make([]int64, n)
	fact[0] = 1

	for i := 1; i < n; i++ {
		fact[i] = fact[i-1] * int64(i) % MOD
	}

	invFact := make([]int64, n)
	invFact[n-1] = modPow(fact[n-1], MOD-2)

	for i := n - 1; i >= 1; i-- {
		invFact[i-1] = invFact[i] * int64(i) % MOD
	}

	comb := fact[n-1]
	comb = comb * invFact[k] % MOD
	comb = comb * invFact[n-1-k] % MOD

	diffCount := n - 1 - k

	ans := int64(m) % MOD
	ans = ans * comb % MOD
	ans = ans * modPow(int64(m-1), int64(diffCount)) % MOD

	return int(ans)
}

The Go implementation mirrors the Python solution closely.

Because multiplication can exceed 32 bit integer limits, all modular arithmetic is performed using int64. Fast exponentiation is implemented manually through binary exponentiation. Arrays of type []int64 store factorials and inverse factorials.

Worked Examples

Example 1

Input:

n = 3
m = 2
k = 1

Compute:

d = n - 1 - k
  = 2 - 1
  = 1

Binomial coefficient:

C(2, 1) = 2

Power term:

(m - 1)^d
= 1^1
= 1

Final answer:

Component	Value
m	2
C(2,1)	2
1^1	1

2 × 2 × 1 = 4

Output:

4

Example 2

Input:

n = 4
m = 2
k = 2

Compute:

d = 3 - 2 = 1

Binomial coefficient:

C(3,2) = 3

Power term:

1^1 = 1

Final answer:

Component	Value
m	2
C(3,2)	3
1^1	1

2 × 3 × 1 = 6

Output:

6

Example 3

Input:

n = 5
m = 2
k = 0

Compute:

d = 4

Binomial coefficient:

C(4,0) = 1

Power term:

1^4 = 1

Final answer:

Component	Value
m	2
C(4,0)	1
1^4	1

2 × 1 × 1 = 2

Output:

2

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n)	Factorial and inverse factorial precomputation dominate
Space	O(n)	Stores factorial and inverse factorial arrays

The combinatorial formula itself is constant time once the factorial tables exist. Building the factorial and inverse factorial arrays requires linear time and linear memory in n, which is easily manageable for n ≤ 100000.

Test Cases

sol = Solution()

assert sol.countGoodArrays(3, 2, 1) == 4      # example 1
assert sol.countGoodArrays(4, 2, 2) == 6      # example 2
assert sol.countGoodArrays(5, 2, 0) == 2      # example 3

assert sol.countGoodArrays(1, 1, 0) == 1      # single element
assert sol.countGoodArrays(1, 100000, 0) == 100000  # any value allowed

assert sol.countGoodArrays(2, 2, 1) == 2      # both elements equal
assert sol.countGoodArrays(2, 2, 0) == 2      # both elements different

assert sol.countGoodArrays(3, 1, 2) == 1      # only one possible array
assert sol.countGoodArrays(3, 1, 0) == 0      # impossible with one value

assert sol.countGoodArrays(4, 3, 3) == 3      # all adjacent equal
assert sol.countGoodArrays(4, 3, 0) == 24     # all adjacent different

assert sol.countGoodArrays(5, 2, 4) == 2      # constant arrays only
assert sol.countGoodArrays(5, 2, 3) == 8      # one differing gap

# stress case, verifies large constraints run efficiently
sol.countGoodArrays(100000, 100000, 50000)

Test Summary

Test	Why
(3,2,1)	Official example 1
(4,2,2)	Official example 2
(5,2,0)	Official example 3
(1,1,0)	Smallest possible input
(1,100000,0)	Single element with maximum value range
(2,2,1)	All pairs equal
(2,2,0)	All pairs different
(3,1,2)	Only one valid array exists
(3,1,0)	Impossible configuration
(4,3,3)	Maximum possible equal pairs
(4,3,0)	No equal pairs
(5,2,4)	Entire array constant
(5,2,3)	Exactly one differing transition
(100000,100000,50000)	Maximum scale performance test

Edge Cases

Array Length One

When n = 1, there are no adjacent pairs at all. Therefore the only valid value of k is 0, which the constraints guarantee. Every choice of the single element produces a valid array, so the answer is simply m. The formula naturally gives:

m × C(0,0) × (m-1)^0 = m

which is correct.

Only One Available Value

When m = 1, every position must contain the value 1. Therefore every adjacent pair is equal.

If:

k = n - 1

there is exactly one valid array.

Otherwise there are zero valid arrays.

The formula handles this automatically because:

(m - 1) = 0

and any required differing gap contributes a factor of zero.

All Adjacent Pairs Equal

When:

k = n - 1

the array must be constant.

There are exactly m possible constant arrays, one for each starting value. The formula becomes:

m × C(n-1,n-1) × (m-1)^0
= m

which matches the direct counting argument.

No Adjacent Pairs Equal

When:

k = 0

every adjacent pair must differ.

The first element has m choices, and each subsequent element has m-1 choices. Therefore:

m × (m-1)^(n-1)

arrays exist.

Since:

C(n-1,0) = 1

the combinatorial formula reduces exactly to this expression.