如何统计满足i<j且a[i]+a[j]可被k整除的数对 寻求低于O(N²)的解法

可以在**O(N + k)**的时间复杂度内解决该问题，远低于O(N²)。

核心原理 #

两个整数a、b的和能被k整除的充要条件是：(a % k + b % k) % k == 0，由此可以按余数分类统计组合数：

• 余数为0的元素，两两组合都满足条件，组合数为 cnt[0] * (cnt[0] - 1) // 2
• 对于满足 1 <= r < k/2 的余数r，每个余数为r的元素都可以和每个余数为k - r的元素配对，总贡献为 cnt[r] * cnt[k - r]
• 如果k是偶数，余数为k/2的元素两两组合也满足条件，组合数为 cnt[k//2] * (cnt[k//2] - 1) // 2

实现步骤 #

1. 初始化一个长度为k的计数数组cnt，所有元素初始值为0
2. 遍历原数组一次，统计每个余数出现的次数：对于每个元素num，计算r = num % k，执行cnt[r] += 1，该步骤时间复杂度O(N)
3. 按照上述规则计算总符合条件的数对数量，仅需要遍历计数数组前半段，该步骤时间复杂度O(k)

参考代码（Python） #

def count_pairs(a, k):
    cnt = [0] * k
    for num in a:
        cnt[num % k] += 1
    res = 0
    # 处理余数为0的配对
    res += cnt[0] * (cnt[0] - 1) // 2
    # 遍历前半段余数
    for r in range(1, (k // 2) + 1):
        if r == k - r:
            # 余数刚好是k的一半，两两配对
            res += cnt[r] * (cnt[r] - 1) // 2
        else:
            # 余数r和k-r配对
            res += cnt[r] * cnt[k - r]
    return res

复杂度说明 #

• 时间复杂度：O(N + k)，仅需要遍历一次原数组和半次计数数组，当k远小于N时，复杂度近似为O(N)
• 空间复杂度：O(k)，仅需要存储长度为k的计数数组

内容的提问来源于stack exchange，提问作者dmitry