互异分划求解子集和整除问题：递归超限与大数计算优化问询

问题背景 #

我们需要计算集合{1, 2, ..., n}中元素和能被5整除的子集数量。3Blue1Brown的Grant Sanderson给出了纯数学公式解法，但直接实现的效率不高；我基于欧拉互异分划的生成函数思路实现时，遇到了两个核心问题：递归深度超限、大数计算效率低下。

现有实现的问题分析 #

1. 递归深度超限：原递归实现的p(k)函数，当n=1000时，递归调用次数远超Python默认的递归深度限制（默认约1000），直接抛出RecursionError。
2. 空间与时间浪费：原递归和迭代实现都维护了完整的子集和计数数组，n=2000时数组长度会达到百万级别，内存占用大，且迭代版时间复杂度为O(n²)，计算效率极低。
3. 大数处理冗余：直接计算所有子集和的计数会产生极大数值，即使使用gmpy2，也因不必要的计算导致效率低下。

优化方案 #

核心思路：状态压缩+迭代更新 #

生成函数的本质是$\prod_{k=1}^n (1 + x^k)$，我们不需要展开整个多项式，只需要跟踪每个模5余数对应的子集数量——因为最终只关心和≡0 mod5的子集数。将状态压缩为长度为5的数组，每个元素dp[r]表示当前子集和模5余r的子集数量，每次处理一个数时，通过平移更新状态（选或不选当前数）。

优化后的代码 #

import gmpy2

def count_subsets_mod5(n):
    # 初始化：dp[r] 记录子集和模5余r的数量，初始仅空集满足余0
    dp = [gmpy2.mpz(0)] * 5
    dp[0] = gmpy2.mpz(1)
    
    for num in range(1, n + 1):
        num_mod = num % 5
        # 复制当前状态，避免更新时覆盖原始值
        new_dp = dp.copy()
        for remainder in range(5):
            # 选择当前数：原余remainder的子集加上num后，余数变为(remainder + num_mod) %5
            new_dp[(remainder + num_mod) % 5] += dp[remainder]
        dp = new_dp
    
    # 返回和能被5整除的子集数量（包含空集，与Grant公式结果一致）
    return dp[0]

def grant_sanderson_formula(n):
    # Grant的数学公式：(2^n + 4*2^{n//5}) //5
    pow_n = gmpy2.mpz(2) ** n
    pow_div5 = gmpy2.mpz(2) ** (n // 5)
    return (pow_n + 4 * pow_div5) // 5

# 测试n=2000的情况
n = 2000
opt_result = count_subsets_mod5(n)
grant_result = grant_sanderson_formula(n)

print("优化后生成函数结果：", opt_result)
print("Grant公式结果：", grant_result)
print("结果一致：", opt_result == grant_result)

优化点说明 #

1. 递归转迭代：完全采用循环实现，彻底解决递归深度超限问题，n取任意大值（如2000、10000）都能正常运行。
2. 空间压缩：仅维护长度为5的数组，空间复杂度O(1)，内存占用可以忽略。
3. 时间高效：时间复杂度O(n)，每个数仅需5次操作，计算速度极快。
4. 大数优化：用gmpy2的mpz类型处理大数，避免Python原生int的性能损耗，同时仅计算必要的模5状态，减少冗余计算。

验证说明 #

优化后的生成函数实现结果与Grant的数学公式完全一致，既保留了生成函数思路的严谨性，又解决了原实现的效率与递归问题。

内容的提问来源于stack exchange，提问作者Darshan P.