Skip to content

2176. Count Equal and Divisible Pairs in an Array 👍

  • Time: $O(n\sqrt{k})$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Solution {
 public:
  int countPairs(vector<int>& nums, int k) {
    int ans = 0;
    unordered_map<int, vector<int>> numToIndices;

    for (int i = 0; i < nums.size(); ++i)
      numToIndices[nums[i]].push_back(i);

    for (const auto& [_, indices] : numToIndices) {
      unordered_map<int, int> gcds;
      for (const int i : indices) {
        const int gcd_i = gcd(i, k);
        for (const auto& [gcd_j, count] : gcds)
          if (gcd_i * gcd_j % k == 0)
            ans += count;
        ++gcds[gcd_i];
      }
    }

    return ans;
  }
};

 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
27
28
class Solution {
  public int countPairs(int[] nums, int k) {
    int ans = 0;
    Map<Integer, List<Integer>> numToIndices = new HashMap<>();

    for (int i = 0; i < nums.length; ++i) {
      numToIndices.putIfAbsent(nums[i], new ArrayList<>());
      numToIndices.get(nums[i]).add(i);
    }

    for (List<Integer> indices : numToIndices.values()) {
      Map<Integer, Integer> gcds = new HashMap<>();
      for (final int i : indices) {
        final int gcd_i = gcd(i, k);
        for (final int gcd_j : gcds.keySet())
          if (gcd_i * gcd_j % k == 0)
            ans += gcds.get(gcd_j);
        gcds.merge(gcd_i, 1, Integer::sum);
      }
    }

    return ans;
  }

  private int gcd(int a, int b) {
    return b == 0 ? a : gcd(b, a % b);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution:
  def countPairs(self, nums: List[int], k: int) -> int:
    ans = 0
    numToIndices = collections.defaultdict(list)

    for i, num in enumerate(nums):
      numToIndices[num].append(i)

    for indices in numToIndices.values():
      gcds = collections.Counter()
      for i in indices:
        gcd_i = math.gcd(i, k)
        for gcd_j, count in gcds.items():
          if gcd_i * gcd_j % k == 0:
            ans += count
        gcds[gcd_i] += 1

    return ans