Skip to content

923. 3Sum With Multiplicity 👍

  • Time:
  • Space: 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
 public:
  int threeSumMulti(vector<int>& A, int target) {
    constexpr int kMod = 1'000'000'007;
    int ans = 0;
    unordered_map<int, long> count;

    for (const int a : A)
      ++count[a];

    for (const auto& [i, x] : count)
      for (const auto& [j, y] : count) {
        int k = target - i - j;
        if (!count.count(k))
          continue;
        if (i == j && j == k)
          ans = (ans + x * (x - 1) * (x - 2) / 6) % kMod;
        else if (i == j && j != k)
          ans = (ans + x * (x - 1) / 2 * count[k]) % kMod;
        else if (i < j && j < k)
          ans = (ans + x * y * count[k]) % kMod;
      }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Solution:
  def threeSumMulti(self, A: List[int], target: int) -> int:
    kMod = 1_000_000_007
    ans = 0
    count = Counter(A)

    for i, x in count.items():
      for j, y in count.items():
        k = target - i - j
        if k not in count:
          continue
        if i == j and j == k:
          ans = (ans + x * (x - 1) * (x - 2) // 6) % kMod
        elif i == j and j != k:
          ans = (ans + x * (x - 1) // 2 * count[k]) % kMod
        elif i < j and j < k:
          ans = (ans + x * y * count[k]) % kMod

    return ans % kMod