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 sum = 0, 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.     	ll pid = pids[i];
  32.         sum = (sum + pid) % k;
  33.  
  34.         // number of subarrays ending here with sum mod k 0 = number of prefixes we can
  35.         // remove with sum mod k = sum
  36.         ans += mp[sum];
  37.         mp[sum]++;
  38.  
  39.     }
  40.  
  41.     return ans;
  42.  
  43. }
  44.  
  45. int main() {
  46. 	ll k = 4;
  47.     vector<ll> pids{1,3,2,4};
  48.  
  49.     cout << f(pids, k) << '\n';
  50.  
  51.     return 0;
  52. }
Success #stdin #stdout 0.01s 5316KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
2
https://ideone.com/Eiysta
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!