fork download
  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
stdin
Standard input is empty
stdout
Case 1 Answer: 3
Case 2 Answer: 5
Case 3 Answer: 0
Case 4 Answer: 0
Case 5 Answer: 3