fork download
copy 
  1. #include <iostream>
  2. #include <vector>
  3. #include <unordered_map>
  4.  
  5. using namespace std;
  6.  
  7. long long countSubarrays(const vector<long long>& nums, int k) {
  8.     // A subarray size must be >= 1. If it must also be < k, 
  9.     // then k must be at least 2 for any valid subarray to exist.
  10.     if (k <= 1) return 0;
  11.  
  12.     unordered_map<int, int> freq;
  13.     vector<int> history; 
  14.  
  15.     long long ans = 0;
  16.     long long prefix_sum = 0;
  17.  
  18.     // Base case: represents the "empty" prefix sum before index 0.
  19.     // Our formula groups the terms as: (Prefix_Sum - Index) % k
  20.     // For the empty prefix: Sum is 0, Index is -1.
  21.     // Key = (0 - (-1)) % k = 1 % k.
  22.     int base_key = 1 % k;
  23.     freq[base_key]++;
  24.     history.push_back(base_key);
  25.  
  26.     for (int i = 0; i < nums.size(); ++i) {
  27.         prefix_sum += nums[i];
  28.  
  29.         // Sliding window logic: 
  30.         // The maximum valid subarray length is k - 1. 
  31.         // If our current index i >= k - 1, the prefix at index (i - k + 1) in our 
  32.         // history array is now too far away and must be removed from the active map.
  33.         if (i >= k - 1) {
  34.             int key_to_remove = history[i - k + 1];
  35.             freq[key_to_remove]--;
  36.         }
  37.  
  38.         // Calculate the current key: (P[i] - i) % k
  39.         long long val = (prefix_sum - i) % k;
  40.  
  41.         // C++ modulo on negative numbers yields a negative result.
  42.         // We wrap it securely to guarantee a positive key.
  43.         if (val < 0) {
  44.             val = (val + k) % k;
  45.         }
  46.  
  47.         int current_key = val;
  48.  
  49.         // If we've seen this key in our active sliding window, it forms a valid subarray
  50.         ans += freq[current_key];
  51.  
  52.         // Record the current key in the map and history for future iterations
  53.         freq[current_key]++;
  54.         history.push_back(current_key);
  55.     }
  56.  
  57.     return ans;
  58. }
  59. int main() {
  60.     vector<pair<vector<long long>, int>> test_cases = {
  61.         {{-2, 1}, 3},
  62.         {{1, 1, 1}, 3},
  63.         {{5, 10, 15}, 1},
  64.         {{4, 4, 4}, 4},
  65.         {{1000000001, 1000000001}, 5}
  66.     };
  67.  
  68.     for (int i = 0; i < test_cases.size(); ++i) {
  69.         long long result = countSubarrays(test_cases[i].first, test_cases[i].second);
  70.         cout << "Case " << i + 1 << " Answer: " << result <<endl;
  71.     }
  72.  
  73.     return 0;
  74. }
Success #stdin #stdout 0s 5324KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Case 1 Answer: 3
Case 2 Answer: 5
Case 3 Answer: 0
Case 4 Answer: 0
Case 5 Answer: 3
https://ideone.com/8m4xZV
language:
C++14 (gcc 8.3)
created:
52 minutes 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!