题目描述

给定正整数K和N，求所有小于等于N的自然数中与K取模的结果之和。

解题思路

我们可以考虑直接遍历1到N的所有自然数，然后对于每个自然数，求其与K取模的结果，最后将所有结果求和即可。时间复杂度为O(N)。

但是，这种方法在N很大的时候会造成很大的计算量，因此我们需要寻找更加高效的算法。

我们可以考虑将所有小于K的数和K本身分别对K取模，可以得到0到K-1这K个不同的取模结果。我们可以发现，对于与K取模的结果相同的自然数，其和与K取模的结果相同。比如，对于K=3来说，所有与3取模为0的自然数之和一定等于所有与3取模为3的自然数之和，因为它们个数相等，且每个数都可以唯一对应到另一个数。

因此，我们可以把1到N的所有自然数分成K个组，每组中的数与K取模的结果相同，然后分别求每组内的所有数之和，最后把这K个结果相加即可。时间复杂度为O(K)。

具体实现可以采用公式(1+2+...+n) = n(n+1)/2来求和，这样可以在常数级别上减少计算量。

代码实现

以下是Python对应的实现代码，时间复杂度为O(K)。

def sum_mod_k(k: int, n: int) -> int:
    res = 0
    # 对于取模结果为0~k-1的k个组，分别计算每组内的自然数之和
    for i in range(k):
        # 每组内有多少个自然数
        cnt = (n - i) // k + 1
        # 每组内所有自然数之和，采用公式(n*(n+1)//2)，可以在常数级别上减少计算量
        s = cnt * (2 * i + (cnt - 1) * k) // 2
        res += s
    return res

参考文献

• 《算法竞赛进阶指南》