aboutsummaryrefslogtreecommitdiff
path: root/count_equal_div_pair_arr.py
blob: 84d4a55ad9afda162757ac0fe85b39b176d6a911 (plain)
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
29
30
31
32
33
34
35
36
37
# 2176. Count Equal and Divisible Pairs in an Array
import collections
from math import gcd

"""
given a 0-indexed integer array 'nums' of length 'n' and an integer 'k',
return the number of pairs '(i, j)' where '0 <= i < j < n' such that 'nums[i]
== nums[j]' and '(i * j)' is divisble by 'k'.
"""


class Solution(object):
    def countPairs(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        ids = collections.defaultdict(list)
        for i, j in enumerate(nums):
            ids[j].append(i)
        ans = 0
        for i in ids.values():
            gcds = collections.Counter()
            for j in i:
                curr_i = gcd(j, k)
                ans += sum(
                    cnt for curr_j, cnt in gcds.items() if curr_i * curr_j % k == 0
                )
                gcds[curr_i] += 1
        return ans


if __name__ == "__main__":
    obj = Solution()
    print(obj.countPairs(nums=[3, 1, 2, 2, 2, 1, 3], k=2))
    print(obj.countPairs(nums=[1, 2, 3, 4], k=1))