Python中含cumsum与迭代切片的for循环如何用numpy向量化实现

解决方案 #

完全可以通过numpy向量化实现，且时间复杂度为O(n)，无需构造高维数组：

import numpy as np

n = 100
a = np.random.random(n)
b = np.random.random(n)
c = np.random.random(n)

# 原循环实现
valid_loop = np.empty(n, dtype=bool)
for i in range(n):
    valid_loop[i] = np.any(a[i] > b[i:] + c[i:].cumsum())

# 向量化实现
pre_c = np.concatenate([[0], np.cumsum(c)])
d = b + pre_c[1:]
min_suffix_d = np.minimum.accumulate(d[::-1])[::-1]
valid_vec = (a + pre_c[:-1]) > min_suffix_d

# 验证结果一致
print(np.all(valid_loop == valid_vec)) # 输出 True

推导逻辑 #

1. 先对c计算前缀和pre_c，头部补0，此时c[i:].cumsum()[k]（k从0开始对应原数组索引i+k）等价于pre_c[i + k + 1] - pre_c[i]
2. 原判断条件a[i] > b[i + k] + c[i:].cumsum()[k]移项后可得：a[i] + pre_c[i] > b[i + k] + pre_c[i + k + 1]
3. 构造数组d[k] = b[k] + pre_c[k+1]，问题转化为：对每个i，是否存在k >= i使得a[i] + pre_c[i] > d[k]
4. 只需要提前计算d的后缀最小值数组，只要a[i] + pre_c[i]大于d从i开始的最小值，就说明存在符合条件的k，和你之前用后缀最小值处理b的思路完全一致

性能说明 #

• 原循环时间复杂度为O(n²)，n较大时（比如n>1000）性能会急剧下降
• 向量化实现所有操作都是O(n)时间复杂度，n=1e5也可以在毫秒级完成计算

内容的提问来源于stack exchange，提问作者trevor-pope