2183. Count Array Pairs Divisible by K

Problem

Given a 0-indexed integer array nums of length n and an integer k, return the number of pairs (i, j) such that:

• 0 <= i < j <= n - 1 and
• nums[i] * nums[j] is divisible by k.

Example 1:

Input: nums = [1,2,3,4,5], k = 2
Output: 7
Explanation:
The 7 pairs of indices whose corresponding products are divisible by 2 are
(0, 1), (0, 3), (1, 2), (1, 3), (1, 4), (2, 3), and (3, 4).
Their products are 2, 4, 6, 8, 10, 12, and 20 respectively.
Other pairs such as (0, 2) and (2, 4) have products 3 and 15 respectively, which are not divisible by 2.

Example 2:

Input: nums = [1,2,3,4], k = 5
Output: 0
Explanation: There does not exist any pair of indices whose corresponding product is divisible by 5.

Constraints:

• 1 <= nums.length <= 10^5
• 1 <= nums[i], k <= 10^5

Problem Summary

Given a 0-indexed integer array nums of length n and an integer k, return the number of index pairs (i, j) that satisfy the following conditions:

• 0 <= i < j <= n - 1 and
• nums[i] * nums[j] is divisible by k.

Solution Ideas

• First, find the greatest common divisor of each element in num with k. Also count the frequency of these gcds and store the data in a map. During the computation, the loop only needs to go up to {{< katex >}}{O(\sqrt {k})}{{< /katex >}}, because every gcd[i] must be a factor of k, and its frequency will not exceed {{< katex >}}{O(\sqrt {k})}{{< /katex >}}. Here is a simple proof: suppose the factors v and k/v are both factors of k. At least one of v and k/v must be less than or equal to {{< katex >}}\sqrt {k}{{< /katex >}}. Therefore, the number of factors of k will not exceed 2 * {{< katex >}}\sqrt {k}{{< /katex >}} = {{< katex >}}{O(\sqrt {k})}{{< /katex >}}.
• After obtaining the above map, use two nested loops to brute-force traverse the key values. If a * b is divisible by k, and a and b are not the same, then the product of the value values corresponding to a and b is the number of index pairs that satisfy the condition; if a and b are the same, then the number of index pairs is {{< katex >}}C_{n}^{2}{{< /katex >}}. Finally, accumulate the result.

Code

package leetcode

import "math"

func countPairs(nums []int, k int) int64 {
	n := int(math.Sqrt(float64(k)))
	gcds, res := make(map[int]int, n), 0
	for _, num := range nums {
		gcds[gcd(num, k)]++
	}

	for a, n1 := range gcds {
		for b, n2 := range gcds {
			if a > b || (a*b)%k != 0 {
				continue
			}
			if a != b {
				res += n1 * n2
			} else { // a == b
				res += n1 * (n1 - 1) / 2
			}
		}
	}
	return int64(res)
}

func gcd(a, b int) int {
	for a%b != 0 {
		a, b = b, a%b
	}
	return b
}