Python递归实现Suffix Sum遇内存超限问题求助

题目背景 #

给定两个数𝑁和𝑀，以及一个包含𝑁个数的数组𝐴。计算最后𝑀个数的和。
注意：必须使用递归解决此问题。
输入：
第一行包含两个数𝑁和𝑀（1 ≤ 𝑀 ≤ 𝑁 ≤ 10⁵）。
第二行包含𝑁个数（−10⁹ ≤ 𝐴ᵢ ≤ 10⁹）。
输出：
输出给定数组最后𝑀个数的和。

问题重现 #

尝试的Python代码：

def suffix_sum(arr, m):
    if m <= 0:
        return 0
    else:
        return arr[-1] + suffix_sum(arr[:-1], m - 1)

n, m = map(int,input().split())
arr = list(map(int,input().split()))
print(suffix_sum(arr, m))

本地样例运行正常，但提交到CodeForces后，测试用例2（推测为10⁴或10⁵级大型数组）出现Memory limit exceeded错误。

内存超限的核心原因 #

• 数组切片的冗余内存消耗：每次递归调用arr[:-1]都会生成一个新的数组副本。当M达到10⁵时，会产生10⁵个长度递减的数组，总内存占用为O(M²)，这是导致内存爆炸的主要原因。
• 递归栈的额外压力：Python默认递归深度约为1000，10⁵级别的递归调用本身也会触发栈溢出，但本题中切片的内存开销是更优先的问题。

递归实现的改进方案 #

1. 传递索引而非数组切片

直接通过索引访问元素，避免创建数组副本，仅传递当前末尾索引和剩余计数：

import sys
sys.setrecursionlimit(1 << 25)  # 提升递归深度限制

def suffix_sum(arr, end_idx, m):
    if m <= 0:
        return 0
    return arr[end_idx] + suffix_sum(arr, end_idx - 1, m - 1)

n, m = map(int, input().split())
arr = list(map(int, input().split()))
print(suffix_sum(arr, n - 1, m))

此方案内存复杂度降至O(M)（仅递归栈占用），彻底解决了切片带来的内存问题。

2. 尾递归模拟（可选）

Python原生不支持尾递归优化，但可以手动用循环模拟尾递归逻辑，进一步降低栈压力：

import sys
sys.setrecursionlimit(1 << 25)

def suffix_sum(arr, m):
    def tail_rec(end_idx, remaining, current_sum):
        if remaining <= 0:
            return current_sum
        return tail_rec(end_idx - 1, remaining - 1, current_sum + arr[end_idx])
    return tail_rec(len(arr)-1, m, 0)

n, m = map(int, input().split())
arr = list(map(int, input().split()))
print(suffix_sum(arr, m))

核心逻辑依然是通过索引访问避免数组副本，同时将递归转化为尾递归形式，减少栈帧的额外开销。

关于Python/PyPy的适配问题 #

PyPy的内存管理和递归性能优于CPython，但原始代码中的数组切片副本问题在PyPy中同样存在，只有修改代码逻辑避免副本，才能通过测试。

内容的提问来源于stack exchange，提问作者Crystal Fire