fork download
  1. #include <iostream>
  2. #include <vector>
  3. #include <algorithm>
  4. using namespace std;
  5. #define int long long
  6. namespace
  7. {
  8.  
  9. int N, M;
  10. vector<int> A;
  11. vector<int> prefixSum;
  12. vector<int> moddedPrefixBIT;
  13. int ans;
  14.  
  15. void inputInfo();
  16.  
  17. void findAns();
  18.  
  19. void addToBit(int x, int value);
  20. int query(int x);
  21.  
  22. void outputAns();
  23. }
  24.  
  25. int32_t main()
  26. {
  27. while (cin >> N >> M)
  28. {
  29. inputInfo();
  30. findAns();
  31. outputAns();
  32. }
  33. }
  34.  
  35. namespace
  36. {
  37. void inputInfo()
  38. {
  39. A.clear(), A.resize(N + 1);
  40. for (int i = 1; i <= N; i++)
  41. {
  42. cin >> A[i];
  43. }
  44. }
  45.  
  46. void findAns()
  47. {
  48. // Mod all values in initial array
  49. for (int i = 1; i <= N; i++)
  50. {
  51. A[i] %= M;
  52. }
  53.  
  54. // Prefix sum
  55. prefixSum.clear(), prefixSum.resize(N + 1);
  56. prefixSum[0] = 0;
  57. for (int i = 1; i <= N; i++)
  58. {
  59. prefixSum[i] = prefixSum[i - 1] + A[i];
  60. }
  61.  
  62. // Calculate total sum
  63. int totalSum = 0;
  64. for (int i = 1; i <= N; i++)
  65. {
  66. totalSum += prefixSum[i] % M;
  67. }
  68.  
  69. // Create Fenwick tree
  70. moddedPrefixBIT.clear(), moddedPrefixBIT.resize(M + 1);
  71. for (int i = 1; i <= N; i++)
  72. {
  73. addToBit(prefixSum[i] % M + 1, 1);
  74. }
  75.  
  76. // Calculate answer
  77. ans = 0;
  78. int reducedSum = 0;
  79. int termsCount = N;
  80. for (int i = 1; i <= N; i++)
  81. {
  82. // Count # of modded prefix sum < reducedSum
  83. int negativeReducedCount = query(reducedSum);
  84.  
  85. // Update answer
  86. ans += totalSum;
  87. ans -= reducedSum * termsCount;
  88. ans += negativeReducedCount * M;
  89.  
  90. // Update state
  91. totalSum -= prefixSum[i] % M;
  92. reducedSum += A[i];
  93. reducedSum %= M;
  94. termsCount--;
  95. addToBit(prefixSum[i] % M + 1, -1);
  96. }
  97. }
  98.  
  99. void addToBit(int x, int value)
  100. {
  101. for (int i = x; i <= M; i += (i & (-i)))
  102. {
  103. moddedPrefixBIT[i] += value;
  104. }
  105. }
  106.  
  107. int query(int x)
  108. {
  109. int value = 0;
  110. for (int i = x; i > 0; i -= (i & (-i)))
  111. {
  112. value += moddedPrefixBIT[i];
  113. }
  114. return value;
  115. }
  116.  
  117. void outputAns()
  118. {
  119. cout << ans << '\n';
  120. }
  121. }
  122.  
Success #stdin #stdout 0.01s 5272KB
stdin
Standard input is empty
stdout
Standard output is empty