fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. using ll = long long;
  5.  
  6. ll f(vector<ll>& pids, ll k) {
  7. 	if (k == 1) return 0;
  8.  
  9. 	int n = pids.size();
  10.     // subtract 1
  11.     for (ll& pid : pids) {
  12.         pid = (pid % k - 1 + k) % k;
  13.     }
  14.  
  15.     vector<ll> pre(n);
  16.     pre[0] = pids[0];
  17.     for (int i = 1; i < n; ++i) {
  18.     	pre[i] += pre[i - 1] + pids[i];
  19.     	pre[i] %= k;
  20.     }
  21.  
  22.     // count subarrays whose sum mod k = 0
  23.     map<ll, ll> mp; // mp[i] = how prefixes have sum mod k = i
  24.     mp[0] = 1; // represents empty prefix
  25.     ll ans = 0;
  26.     for (int i = 0; i < n; ++i) {
  27.     	// remove prefixes that leave suffix len >= k
  28.     	if (i == k - 1) mp[0]--;
  29.     	if (i >= k) mp[pre[i-k]]--;
  30.  
  31.         // number of subarrays ending here with sum mod k 0 = number of prefixes we can
  32.         // remove with sum mod k = pre[i]
  33.         ans += mp[pre[i]];
  34.         mp[pre[i]]++;
  35.  
  36.     }
  37.  
  38.     return ans;
  39.  
  40. }
  41.  
  42. int main() {
  43. 	ll k = 4;
  44.     vector<ll> pids{1,3,2,4};
  45.  
  46.     cout << f(pids, k) << '\n';
  47.  
  48.     return 0;
  49. }
Success #stdin #stdout 0.01s 5288KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
2
https://ideone.com/RPCPtY
language:
C++14 (gcc 8.3)
created:
5 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!