給定一串?dāng)?shù)nums和整數(shù)k,問這串?dāng)?shù)的子串和可以被k整除的子串(連續(xù)的一段)有多少
參考了Discussion:
先預(yù)處理prefix_sum(%k之后的值)
如果子串nums[i]~num[j]之和可以被k整除,說(shuō)明prefix_sum[i] = prefix_sum[j],于是用dict存各種prefix_sum的可能性有多少,最后
ans = sum(prefix_sum[x] * (prefix_sum[x]- 1) / 2), x=0~k-1
因?yàn)槭荂(n, 2)的組合數(shù),所以不必等所有prefix_sum計(jì)算完再一個(gè)個(gè)算,在掃描nums,計(jì)算當(dāng)前prefix_sum的時(shí)候順便累加即可
注意初始化prefix_sum[0] = 1,因?yàn)橐粋€(gè)數(shù)都不取的話和為0
1 #974
2 #Runtime: 232 ms (Beats 92.68%)
3 #Memory: 16.8 MB (Beats 37.80%)
4
5 class Solution(object):
6 def subarraysDivByK(self, nums, k):
7 """
8 :type nums: List[int]
9 :type k: int
10 :rtype: int
11 """
12 pre_sum = defaultdict(int)
13 t, ans = 0, 0
14 pre_sum[0] = 1
15 for i in nums:
16 t = (t + i) % k
17 pre_sum[t] += 1
18 ans += pre_sum[t] - 1
19 return ans