题目是leetcode0974 和可被k整除的子数组
给定一个数组,以及目标k,求所有和为k的倍数的连续子序列
输入:
nums = [45, 2, 3, 9, 1, 5]
k = 6
输出:
[[3, 9], [45, 2, 3, 9, 1], [3, 9, 1, 5], [1, 5]]
类似前缀和的想法,考虑前i项和modk等于0, 1, ..., k-1的和,模k同余的项之间的连续子序列之和为k的倍数。可以用一个字典来记住这些模k同余的项。
import collections
def solution(A, k):
n = len(A)
summ = 0
ans = []
sol_map = collections.defaultdict(list)
sol_map[0] = [0] # 空子序列的和为0
for i in range(1, n+1): # 从第1个数到第n个数
summ = (A[i-1] + summ) % k
if sol_map[summ]:
for idx in sol_map[summ]:
ans.append(A[idx: i])
sol_map[summ].append(i)
return ans
如果只是想数个数,字典改为模k同余的项的个数即可
import collections
def solution2(A, k):
n = len(A)
summ = 0
ans = 0
sol_map = collections.defaultdict(int)
sol_map[0] = 1 # 空子序列的和为0
for i in range(1, n+1):
summ = (A[i-1] + summ) % k
if sol_map[summ]:
ans += sol_map[summ]
sol_map[summ] += 1
return ans