统计可以被 K 整除的下标对数目

标签: 数组 数学 数论

难度: Hard

给你一个下标从 0 开始、长度为 n 的整数数组 nums 和一个整数 k ,返回满足下述条件的下标对 (i, j) 的数目:

  • 0 <= i < j <= n - 1
  • nums[i] * nums[j] 能被 k 整除。

示例 1:

输入:nums = [1,2,3,4,5], k = 2
输出:7
解释:
共有 7 对下标的对应积可以被 2 整除:
(0, 1)、(0, 3)、(1, 2)、(1, 3)、(1, 4)、(2, 3) 和 (3, 4)
它们的积分别是 2、4、6、8、10、12 和 20 。
其他下标对,例如 (0, 2) 和 (2, 4) 的乘积分别是 3 和 15 ,都无法被 2 整除。    

示例 2:

输入:nums = [1,2,3,4], k = 5
输出:0
解释:不存在对应积可以被 5 整除的下标对。

提示:

  • 1 <= nums.length <= 105
  • 1 <= nums[i], k <= 105

Submission

运行时间: 95 ms

内存: 27.1 MB

class Solution:
    def countPairs(self, nums: List[int], k: int) -> int:
        freq = Counter(gcd(num, k) for num in nums)
        c = sorted(freq.items())
        n = len(c)

        ret = 0
        for i in range(n):
            m = bisect.bisect_left(c, (math.ceil(k / c[i][0]), 0), i + 1)
            for j in range(m, n):
                if c[i][0] * c[j][0] % k == 0:
                    ret += c[i][1] * c[j][1]
            if c[i][0] * c[i][0] % k == 0:
                ret += c[i][1] * (c[i][1] - 1) // 2
        return ret
                

Explain

本题解采用了数学和哈希表的方法来优化解决问题。首先用哈希表freq统计数组中每个元素与k的最大公因数(gcd)出现的频次。然后对这个哈希表中的元素(gcd值及其出现次数)进行排序。对于每一对gcd值,如果它们的乘积能被k整除,则表示存在多个满足条件的下标对(i, j),即这两个gcd值对应的原数组元素组合的所有可能。对于相同的gcd值,需要使用组合数公式来计算可能的下标对数。通过这种方法,将问题转化为了较少数量的gcd值的组合问题,从而避免了直接的n^2复杂度的暴力解法。

时间复杂度: O(u^2 + u log u + n log k)

空间复杂度: O(u)

class Solution:
    def countPairs(self, nums: List[int], k: int) -> int:
        from collections import Counter
        from math import gcd, ceil
        import bisect

        # 计算每个数字和k的最大公因数,并统计每个gcd的出现频率
        freq = Counter(gcd(num, k) for num in nums)
        # 将gcd及其频率按字典序排序
        c = sorted(freq.items())
        n = len(c)

        ret = 0
        # 遍历所有可能的gcd组合
        for i in range(n):
            # 找到第一个可能与c[i]组合满足条件的gcd位置
            m = bisect.bisect_left(c, (ceil(k / c[i][0]), 0), i + 1)
            # 遍历所有可能的组合
            for j in range(m, n):
                if c[i][0] * c[j][0] % k == 0:
                    ret += c[i][1] * c[j][1]
            # 单独处理相同gcd的组合情况
            if c[i][0] * c[i][0] % k == 0:
                ret += c[i][1] * (c[i][1] - 1) // 2
        return ret

Explore

直接考虑 nums[i] 和 nums[j] 的乘积会导致时间复杂度过高,特别是当数组长度较大时。使用最大公因数(gcd)的方法可以将问题简化。如果 nums[i] 和 nums[j] 的乘积能被 k 整除,那么 nums[i] 和 nums[j] 与 k 的 gcd 乘积也一定能被 k 整除。这样通过先计算每个数与 k 的 gcd,再处理这些 gcd 的乘积,可以有效减少需要处理的数据量和复杂度。

bisect_left 函数在 sorted list 中用于找到插入位置以维护列表顺序。它能快速找到第一个大于或等于指定元素的位置。在这个题解中,使用 bisect_left 是为了快速找到能与当前 gcd 值相乘并能被 k 整除的另一个 gcd 的起始位置。它确保能找到正确的位置,因为数组已经被排序,bisect_left 依赖于这个顺序来进行二分查找,从而快速定位到正确的开始位置。

当两个下标 i 和 j 指向的元素具有相同的 gcd 值时,任选两个不同的下标组成的对都满足条件。因此,问题可以转化为如何从频次为 c[i][1] 的集合中选择两个不同元素的组合数。这可以通过组合数公式 C(n, 2) = n*(n-1)/2 来计算,其中 n 是具有相同 gcd 的元素的数量。

排序操作主要影响算法性能,因为它具有 O(N log N) 的时间复杂度,其中 N 是不同 gcd 值的数量。排序是必须的,因为它保证了后续使用二分查找(bisect_left)的有效性和正确性。通过排序,可以确保在查找合适的 gcd 组合时,数组是有序的,这是二分查找算法的前提条件。排序后,可以有效地通过索引快速访问和比较元素,从而提高算法的整体效率和性能。