计算可被K整除的子阵列

Question

给定n个正整数的序列,我们需要计算其总和可被k整除的连续子序列.

约束:N高达10 ^ 6,每个元素高达10 ^ 9,K高达100

示例:设N = 5且K = 3,阵列为1 2 3 4 1

这里的答案是4

说明:存在4个子序列,其总和可以被3整除,它们是

3
1 2
1 2 3
2 3 4

我的尝试:

long long int count=0;
for(int i=0;i<n;i++){
    long long int sum=0;
    for(int j=i;j<n;j++)
    {
        sum=sum+arr[j];
        if(sum%k==0)
        {
            count++;
        }
    }
}

但显然它的方法很差.他们可以更好地解决这个问题吗？请帮忙.

完整问题:https://www.hackerrank.com/contests/w6/challenges/consecutive-subsequences

Answer 1

这是一个快速的O(n + k)解决方案:

1)让计算前缀和pref [i](对于0 <= i <n).

2)现在我们可以计算count [i] - 和i的模数为k的前缀数(0 <= i <k).这可以通过迭代所有前缀并进行计数[pref [i]%k] ++来完成.最初,count [0] = 1(空前缀有0),0表示i!= 0.

3)答案是所有i的总和[i]*(count [i] -1)/ 2.

4)最好以模k计算前缀和以避免溢出.

它为什么有效？让我们仔细看看一个可被k整除的子阵列.假设它从L位置开始并以R位置结束.当且仅当pref [L - 1] == pref [R](模k)时,它可以被k整除,因为它们的差值是零模k(通过可分性的定义).因此,对于每个固定模,我们可以选择带有此前缀sum modulo k的任何两个前缀(并且确实有count [i]*(count [i] - 1)/ 2种方法).

这是我的代码:

long long get_count(const vector<int>& vec, int k) {
  //Initialize count array.
  vector<int> cnt_mod(k, 0);
  cnt_mod[0] = 1;
  int pref_sum = 0;
  //Iterate over the input sequence.
  for (int elem : vec) {
    pref_sum += elem;
    pref_sum %= k;
    cnt_mod[pref_sum]++;
  }
  //Compute the answer.
  long long res = 0;
  for (int mod = 0; mod < k; mod++)
    res += (long long)cnt_mod[mod] * (cnt_mod[mod] - 1) / 2;
  return res;
}