求数组中和可被k整除的数对数量：解法遇性能问题求助

问题描述 #

给定整数数组a和整数k，计算满足i<j且a[i]+a[j]能被k整除的数对数量，约束条件：

• 1≤a.length≤1e5
• 1≤a[i]≤1e9
• 1≤k≤1e9

示例输入：a=[1,2,3,4,5]、k=3，答案为4。

原解法问题分析 #

你原代码使用长度为k的数组作为计数桶，这种方案在k取值极大（比如1e9）时会触发致命问题：

• Java中内存溢出：无法分配长度为1e9的数组，直接抛出OutOfMemoryError。
• Python中超时：创建包含1e9个元素的列表需要占用巨量内存，且初始化过程耗时极长，导致超出时间限制。

优化思路 #

改用哈希字典（Python的dict、Java的HashMap）存储每个模值的出现次数，空间复杂度从O(k)降至O(n)（最多存储n个不同的模值，n为数组长度），同时保持O(n)的时间复杂度。

优化后的代码 #

Python实现 #

def solution(a, k):
    mod_counts = {}
    res = 0
    for num in a:
        mod = num % k
        # 计算当前数需要的互补模值
        complement = (k - mod) % k
        # 累加已有互补模值的计数
        res += mod_counts.get(complement, 0)
        # 更新当前模值的计数
        mod_counts[mod] = mod_counts.get(mod, 0) + 1
    return res

Java实现 #

import java.util.HashMap;
import java.util.Map;

public class Solution {
    public long solution(int[] a, int k) {
        Map<Integer, Integer> modCounts = new HashMap<>();
        long res = 0;
        for (int num : a) {
            int mod = num % k;
            // 处理Java负数取模的特殊情况，转为正模值
            if (mod < 0) mod += k;
            int complement = (k - mod) % k;
            res += modCounts.getOrDefault(complement, 0);
            modCounts.put(mod, modCounts.getOrDefault(mod, 0) + 1);
        }
        return res;
    }
}

关键注意事项 #

1. Java负数模处理：Java中负数取模结果为负数（如(-1) % 3 = -1），需要调整为正模值，否则互补模值计算错误。
2. 结果溢出问题：当数组长度为1e5时，最多可能有1e5*(1e5-1)/2 ≈5e9个数对，超过int的范围（Java中int最大约2e9），因此需用long存储结果。

示例验证 #

以a=[1,2,3,4,5]、k=3为例：

• 遍历1：模为1，互补模为2，无匹配，结果0，字典存{1:1}
• 遍历2：模为2，互补模为1，累加1，结果1，字典存{1:1,2:1}
• 遍历3：模为0，互补模为0，无匹配，结果1，字典存{1:1,2:1,0:1}
• 遍历4：模为1，互补模为2，累加1，结果2，字典存{1:2,2:1,0:1}
• 遍历5：模为2，互补模为1，累加2，结果4，符合预期。

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