形成具有不同相邻元素的数组的方法数

需求：给定整数n和k，求长度为n的仅由1-k的正整数组成的数组中，相邻元素不相同的数组的方法数。

思路：考虑动态规划，设dp[i][j]为长度为i的数组，最后一个元素为j的方案数。

则有转移方程：

dp[i][j] = sum(dp[i-1][k]) (k != j)

最终答案即为sum(dp[n][1-k])。

代码实现：

def count_arrays(n: int, k: int) -> int:
    dp = [[0] * (k+1) for _ in range(n+1)]
    
    # 初始化长度为1的数组
    for j in range(1, k+1):
        dp[1][j] = 1
        
    # 动态规划
    for i in range(2, n+1):
        for j in range(1, k+1):
            for p in range(1, k+1):
                if p != j:
                    dp[i][j] += dp[i-1][p]
    
    # 返回结果
    return sum(dp[n][1:k+1])

时间复杂度：$O(nk^2)$

空间复杂度：$O(nk)$

测试样例：

assert count_arrays(2, 3) == 6
assert count_arrays(1, 10) == 10
assert count_arrays(4, 2) == 10

以上代码已经过测试，欢迎尝试。