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

Difficulty: 🔴 Hard
Topics: Math, Combinatorics

Solution

Problem Understanding

This problem asks us to count how many arrays of length n can be formed using integers between 1 and m such that exactly k adjacent pairs are equal. In other words, we are looking for sequences where the number of times arr[i] equals arr[i-1] is precisely k. The output is the total count of these arrays modulo 10^9 + 7.

The input integers n, m, and k define the size of the array, the range of values allowed in the array, and the number of adjacent matches, respectively. The constraints indicate that n and m can be as large as 10^5, which prevents any brute-force generation of all possible arrays since there are m^n possible arrays. Edge cases include k = 0 (no repeated adjacent elements) and k = n - 1 (all elements equal). Another subtle edge case is m = 1, which restricts all elements to the same value.

The challenge is combinatorial in nature, requiring careful counting rather than enumeration.

Approaches

A brute-force approach would generate all m^n arrays and check for exactly k adjacent equalities. While correct in principle, this is infeasible for the given constraints because the number of arrays grows exponentially.

The optimal approach uses dynamic programming and combinatorics. We can define a recurrence for counting arrays based on their length i and the number of matches j. If the previous element matches, the new element contributes to j + 1; if it differs, j remains unchanged. By carefully counting the number of ways to continue sequences and applying modular arithmetic, we can compute the total efficiently. The key insight is to observe that every array can be partitioned into "blocks" of identical numbers separated by changes, and the number of changes plus one equals the number of blocks. This allows us to use combinatorial coefficients to directly compute counts.

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(m^n)	O(n)	Generate all arrays and count matches. Not feasible for large n or m.
Optimal	O(n * k)	O(k)	Dynamic programming with combinatorial counting using modular arithmetic.

Algorithm Walkthrough

Define a 1D DP array dp[j] representing the number of sequences of length i that have exactly j matches.
Initialize dp[0] = m because for a sequence of length 1, no matches exist and there are m choices.
Iterate through the positions of the array from 2 to n. For each position:

Initialize a new DP array new_dp.
For each j from 0 to k:
Add dp[j] * (m - 1) to new_dp[j] representing sequences where the next element differs from the previous one.
If j > 0, add dp[j - 1] * 1 to new_dp[j] representing sequences where the next element matches the previous one.
Apply modulo operation at each step to avoid overflow.
Replace dp with new_dp.

After processing all n positions, dp[k] contains the total number of arrays with exactly k matching adjacent elements.
Return dp[k] % MOD where MOD = 10^9 + 7.

Why it works: The algorithm tracks all valid sequences incrementally and ensures that the count of matching adjacent pairs is precise at each step. The DP invariant is that dp[j] always represents the number of sequences of the current length with exactly j matches, guaranteeing correctness.

Python Solution

class Solution:
    def countGoodArrays(self, n: int, m: int, k: int) -> int:
        MOD = 10**9 + 7
        dp = [0] * (k + 1)
        dp[0] = m  # Base case: first element has no matches

        for _ in range(1, n):
            new_dp = [0] * (k + 1)
            for j in range(k + 1):
                # Case 1: next element is different
                new_dp[j] += dp[j] * (m - 1)
                new_dp[j] %= MOD
                # Case 2: next element is same (increase match count)
                if j > 0:
                    new_dp[j] += dp[j - 1]
                    new_dp[j] %= MOD
            dp = new_dp

        return dp[k]

The Python implementation closely follows the algorithm. dp stores the current counts of sequences with exactly j matches. At each new position in the array, we compute how these counts evolve if the next element is the same or different. Modulo operations prevent integer overflow.

Go Solution

func countGoodArrays(n int, m int, k int) int {
    MOD := int(1e9 + 7)
    dp := make([]int, k+1)
    dp[0] = m

    for i := 1; i < n; i++ {
        newDp := make([]int, k+1)
        for j := 0; j <= k; j++ {
            newDp[j] = (newDp[j] + dp[j]*(m-1)) % MOD
            if j > 0 {
                newDp[j] = (newDp[j] + dp[j-1]) % MOD
            }
        }
        dp = newDp
    }

    return dp[k]
}

In Go, we use slices instead of arrays and handle integer arithmetic carefully to prevent overflow. The modulo operation is applied immediately after every update.

Worked Examples

Example 1: n = 3, m = 2, k = 1

Initial dp = [2, 0]

Iteration 1 (i = 2):

new_dp[0] = dp[0]*(m-1) = 2*1 = 2
new_dp[1] = dp[0] = 2

After iteration: dp = [2, 2]

Iteration 2 (i = 3):

new_dp[0] = dp[0]*(m-1) = 2*1 = 2
new_dp[1] = dp[0] + dp[1]*(m-1) = 2 + 2*1 = 4

After iteration: dp = [2, 4]

Return dp[1] = 4.

Example 2: n = 4, m = 2, k = 2

Initial dp = [2, 0, 0]

Iteration 1 (i = 2): dp = [2, 2, 0]

Iteration 2 (i = 3): dp = [2, 4, 2]

Iteration 3 (i = 4): dp = [2, 6, 6]

Return dp[2] = 6.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n * k)	We iterate over n positions and for each position update k + 1 states.
Space	O(k)	We only maintain two arrays of size k + 1.

The DP optimization allows us to compute the solution efficiently without generating all sequences.

Test Cases

# Provided examples
assert Solution().countGoodArrays(3, 2, 1) == 4  # Example 1
assert Solution().countGoodArrays(4, 2, 2) == 6  # Example 2
assert Solution().countGoodArrays(5, 2, 0) == 2  # Example 3

# Edge cases
assert Solution().countGoodArrays(1, 1, 0) == 1  # Single element
assert Solution().countGoodArrays(3, 1, 2) == 1  # All same elements
assert Solution().countGoodArrays(5, 3, 0) == 48  # No matches
assert Solution().countGoodArrays(5, 3, 4) == 3   # All elements match except one

Test	Why
Single element	Base case, no adjacent comparisons
All elements same	k = n-1, ensures all equalities handled
No matches	Ensures algorithm correctly counts sequences with alternating elements
Maximum matches	Ensures sequences with mostly identical elements are counted

Edge Cases

One important edge case is when m = 1. In this scenario, all elements must be the same. The only valid k is n - 1. Any other k should return 0. Another edge case is k = 0, which requires all adjacent elements to differ. If m = 1, this is impossible. Finally, n = 1 is special because there are no adjacent pairs. The algorithm correctly handles it by initializing dp[0] = m and returning dp[k], which matches the expected result for all allowed k.