leetcode刷题笔记-prefix sum

1685. Sum of Absolute Differences in a Sorted Array

class Solution:
    def getSumAbsoluteDifferences(self, nums: List[int]) -> List[int]:
#         res[i] = (nums[i] - nums[0]) + (nums[i] - nums[1]) + ... + (nums[i] - nums[i - 1])        <--- absolute difference of nums[i] with first i numbers
#      + (nums[i] - nums[i]) + (nums[i + 1] - nums[i]) + (nums[i + 2] - nums[i]) + ... + (nums[n - 1] - nums[i])   <--- absolute difference of nums[i] with last n - i numbers

#         res[i] = i * nums[i] - (nums[0] + ... + nums[i - 1])        <--- absolute difference of nums[i] with first i numbers
#      + (nums[i + 1] + ... + nums[n]) - (n - i) * nums[i]   <--- absolute difference of nums[i] with last n - i numbers

#      res[i] = i * nums[i] - prefixSum[i] +
#       prefixSum[n] - prefixSum[i] - (n - i) * nums[i]

        pre_s = [0]
        for n in nums:
            pre_s.append(n + pre_s[-1])
        re, n = [], len(nums)
        for i, v in enumerate(nums):
            re.append(i*v - pre_s[i] + pre_s[n] - pre_s[i] - (n-i) *v)
        return re

974. Subarray Sums Divisible by K

class Solution:
    def subarraysDivByK(self, nums: List[int], k: int) -> int:
        '''
        how many ways can we pair up all prefix sum pairs (i,j) where i < j such that (prefix[j] - prefix[i]) % K == 0. we easily get a O(n^2) solution.
        To linear: 
        I noted above that we need to find all prefix sum pairs (i,j) such tha (p[j] - p[i]) % K == 0.
        But this is only true, if and only if p[j] % K == p[i] % K
        '''
 
        count = [0] * k
        count[0] = 1 # [0] can be diviable by k
        preS = re = 0
        for n in nums:
            preS = (preS + n) % k
            re += count[preS]
            count[preS] += 1
        return re