Leetcode 959 Solution

This article provides solution to leetcode question 959 (3sum-with-multiplicity)

https://leetcode.com/problems/3sum-with-multiplicity

Solution

class Solution(object):
    def threeSumMulti(self, A, target):
        """
        :type A: List[int]
        :type target: int
        :rtype: int
        """
        ans = 0

        m = collections.defaultdict(int)
        for a in A:
            m[a] += 1

        keys = m.keys()
        for i in keys:
            for j in keys:
                if i > j:
                    continue
                k = target - i - j

                if i == j == k:
                    ans += m[k] * (m[k] - 1) * (m[k] - 2) / 6
                elif i == j != k:
                    ans += m[i] * (m[i] - 1) * m[k] / 2
                elif i < j < k:
                    ans += m[i] * m[j] * m[k]

        return ans % 1000000007