国际空间研究组织 | ISRO CS 2009 |问题 43

题目描述

给定一个由整数组成的数组arr和一个整数k，找出数组中任意两个元素的和能够被k整除的所有不同的二元组的数量。

函数签名

def count_pairs(arr: List[int], k: int) -> int:
    pass

示例

assert count_pairs([1, 2, 3, 4, 5], 2) == 4
assert count_pairs([5, 9, 36, 74, 52, 31, 42], 3) == 5

题目解析

如果对每两个元素进行暴力判断，时间复杂度会超出限制，而使用哈希表可以使时间复杂度为 O(n)。

建立一个用于存放哈希值的字典哈希表，如果哈希表中存在该元素的模，则将该元素的数量累加到该模值对应的哈希桶中。同时遍历哈希表，当哈希值x对应的哈希桶中元素数量为y时，可以从哈希表中取出 x 和 k-x 的数量相乘再加到答案上。

代码实现

def count_pairs(arr: List[int], k: int) -> int:
    count = 0
    hash_map = {}
    for num in arr:
        mod = num % k
        if mod in hash_map:
            hash_map[mod] += 1
        else:
            hash_map[mod] = 1
    for mod in hash_map.keys():
        if mod == 0:
            count += (hash_map[mod] * (hash_map[mod] - 1)) // 2
        elif mod <= k // 2:
            if k - mod in hash_map:
                count += hash_map[mod] * hash_map[k - mod]
    return count

上面这段代码实现了如何统计符合条件的二元组数量。传入参数为数组 arr 和整数 k，返回值为符合条件的二元组数量的整数值。

该函数先遍历一遍数组 arr，创建一个哈希表 hash_map，用于存储模值对应的哈希桶，增加哈希桶中元素数量。接着遍历哈希表，如果模为0，则将计算结果加上n(n - 1) // 2；如果原数组中存在该模值x，同时又存在模值为k-x的数，则将计算结果加上哈希表中x和k-x对应的桶里元素数量的乘积。最后返回计算结果即可。

测试结果

assert count_pairs([1, 2, 3, 4, 5], 2) == 4
assert count_pairs([5, 9, 36, 74, 52, 31, 42], 3) == 5