数组分区计数与输出方案Python实现：子数组和可被K整除长度2-4

问题说明 #

合法分区需要满足两个核心条件：

• 每个分区得到的子数组元素之和可被K整除
• 每个子数组的长度取值范围为2 ≤ 长度 ≤4
  输入约束：
• 数组长度N满足 0 ≤ N ≤ 10^5
• 数组元素A[i]满足 0 ≤ A[i] ≤ 10^9
• 整数K满足 1 ≤ K ≤ 10^9

实现思路 #

注意：当数组长度达到10^5量级时，合法分区方案数可能呈指数级增长，无法全部存储输出，因此提供两类实现：

1. 仅计数版本：时间复杂度O(N)，空间复杂度O(N)，支持最长10^5长度的数组输入
2. 全方案输出版本：仅适合短数组使用，通过回溯法枚举所有合法分区结果

核心优化逻辑：
子数组和能被K整除等价于该子数组首尾对应的前缀和模K结果相等，避免每次计算子数组和的开销。
计数逻辑基于动态规划：dp[i]表示前i个元素的合法分区方案数，dp[i]等于满足长度2、3、4条件的前序合法方案数之和，边界条件dp[0]=1。

代码实现 #

# 版本1：仅计算合法分区总方案数，支持大数组
def count_valid_partitions(A, K):
    n = len(A)
    if n == 0:
        return 1
    # 前缀和模K预处理
    prefix_mod = [0] * (n + 1)
    for i in range(n):
        prefix_mod[i+1] = (prefix_mod[i] + A[i]) % K
    dp = [0] * (n + 1)
    dp[0] = 1
    for i in range(2, n+1):
        # 检查长度为2的子数组
        if i >= 2 and prefix_mod[i] == prefix_mod[i-2]:
            dp[i] += dp[i-2]
        # 检查长度为3的子数组
        if i >= 3 and prefix_mod[i] == prefix_mod[i-3]:
            dp[i] += dp[i-3]
        # 检查长度为4的子数组
        if i >= 4 and prefix_mod[i] == prefix_mod[i-4]:
            dp[i] += dp[i-4]
    return dp[n]

# 版本2：输出所有合法分区方案，仅适合短数组
def get_all_valid_partitions(A, K):
    n = len(A)
    res = []
    prefix_mod = [0] * (n + 1)
    for i in range(n):
        prefix_mod[i+1] = (prefix_mod[i] + A[i]) % K
    
    def backtrack(start, cur_path):
        if start == n:
            res.append(cur_path.copy())
            return
        # 尝试长度2、3、4的合法子数组
        for length in [2,3,4]:
            end = start + length
            if end > n:
                break
            if prefix_mod[end] == prefix_mod[start]:
                cur_path.append(A[start:end])
                backtrack(end, cur_path)
                cur_path.pop()
    
    backtrack(0, [])
    return res

# 示例测试
if __name__ == "__main__":
    A = [6,3,3,8,4]
    K = 3
    total_count = count_valid_partitions(A, K)
    all_partitions = get_all_valid_partitions(A, K)
    print(f"共{total_count}种方案")
    print("分区结果：", all_partitions)

示例输出 #

运行上述测试代码输出结果：

共1种方案
分区结果： [[[6, 3], [3, 8, 4]]]

内容的提问来源于stack exchange，提问作者Ashutosh Ranjan