923. 3Sum With Multiplicity - cocoder39/coco39_LC GitHub Wiki

923. 3Sum With Multiplicity

class Solution:
    def threeSumMulti(self, arr: List[int], target: int) -> int:
        arr.sort()
        cnt = 0
        MOD = 10**9 + 7
        for i, num in enumerate(arr):
            two_sum = target - num
            cnt += self.twoSum(arr, i+1, len(arr)-1, two_sum)
            cnt %= MOD 
        return cnt
    
    def twoSum(self, arr: List[int], begin: int, end: int, target: int) -> int:
        cnt = 0
        while begin < end:
            if arr[begin] + arr[end] < target:
                begin += 1
            elif arr[begin] + arr[end] > target:
                end -= 1
            elif arr[begin] != arr[end]:
                left, right = 1, 1
                while begin + 1 < end and arr[begin+1] == arr[begin]:
                    begin += 1
                    left += 1
                while end - 1 > begin and arr[end-1] == arr[end]:
                    end -= 1
                    right += 1
                cnt += left * right
                begin += 1
                end -= 1
            else: # arr[begin] + arr[end] == target and arr[begin] == arr[end]
                cnt += (end-begin+1) * (end-begin) // 2
                break
        return cnt