Advertisement
Guest User

Untitled

a guest
Dec 29th, 2019
1,032
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.73 KB | None | 0 0
  1. #include <iostream>
  2. #include <vector>
  3. #include <map>
  4. #include <set>
  5. #include <queue>
  6. #include <algorithm>
  7. #include <string>
  8. #include <cmath>
  9. #include <cstdio>
  10. #include <iomanip>
  11. #include <fstream>
  12. #include <cassert>
  13. #include <cstring>
  14. #include <unordered_set>
  15. #include <unordered_map>
  16. #include <numeric>
  17. #include <ctime>
  18. #include <bitset>
  19. #include <complex>
  20. #include <random>
  21.  
  22. using namespace std;
  23.  
  24. #define int long long
  25.  
  26. int stupid(string s) {
  27.   int n = (int)s.size();
  28.     vector<int> ps(n + 1);
  29.     for (int i = 0; i < n; i++) {
  30.         ps[i + 1] = ps[i] + s[i] - '0';
  31.     }
  32.     int ans = 0;
  33.     for (int i = 0; i < n; i++) {
  34.         for (int j = i + 1; j <= n; j++) {
  35.             if (ps[i] != ps[j] && (j - i) % (ps[j] - ps[i]) == 0) ans++;
  36.         }
  37.     }
  38.     return ans;
  39. }
  40.  
  41. int solve(string s) {
  42.   const int LIM = 300;
  43.  
  44.   int ans = 0;
  45.  
  46.   int n = (int)s.size();
  47.     vector<int> ps(n + 1);
  48.     for (int i = 0; i < n; i++) {
  49.         ps[i + 1] = ps[i] + s[i] - '0';
  50.     }
  51.     vector<int> val(n + 1);
  52.     vector<int> tv(n + 1);
  53.     iota(val.begin(), val.end(), 0);
  54.    
  55.     for (int i = 1; i < LIM; i++) {
  56.       for (int j = 0; j <= n; j++) {
  57.         val[j] -= ps[j];
  58.         tv[j] = val[j];
  59.       }
  60.       sort(tv.begin(), tv.end());
  61.       for (int j = 0; j <= n; j++) {
  62.         int r = j;
  63.         while (r + 1 <= n && tv[r] == tv[r + 1]) {
  64.           r++;
  65.         }
  66.         int len = r - j + 1;
  67.         ans += (len) * (len - 1) / 2;
  68.         j = r;
  69.       }
  70.     }
  71.  
  72.     vector<int> pos;
  73.     for (int i = 0; i < n; i++) {
  74.       if (s[i] == '1') {
  75.         pos.push_back(i);
  76.       }
  77.     }
  78.     pos.push_back(n);
  79.     int possz = (int)pos.size() - 1;
  80.  
  81.     int cur_cnt = 0;
  82.  
  83.     for (int i = 0; i < n; i++) {
  84.       for (int j = cur_cnt; j < min(possz, cur_cnt + n / LIM + 1); j++) {
  85.         int cnt_1 = j - cur_cnt + 1;
  86.         int cur_lol = i + cnt_1 * LIM;
  87.         {
  88.           int delta = pos[j] - cur_lol + 1;
  89.           if (delta > 0) {
  90.             cur_lol += (cnt_1) * ((delta + cnt_1 - 1) / cnt_1);
  91.           }
  92.         }
  93.         if (cur_lol <= pos[j + 1]){
  94.           int delta = pos[j + 1] - cur_lol;
  95.           ans += (delta / cnt_1) + 1;
  96.         }
  97.       }
  98.       if (s[i] == '1') {
  99.         cur_cnt++;
  100.       }
  101.     }
  102.  
  103.     return ans;
  104. }
  105.  
  106. void gen(int n) {
  107.   string s;
  108.   for (int i = 0; i < n; i++) {
  109.     s.push_back('0' + rand() % 2);
  110.   }
  111.   if (stupid(s) != solve(s)) {
  112.     cout << "WA " << s << endl;
  113.     exit(0);
  114.   } else {
  115.     cout << "OK" << endl;
  116.   }
  117. }
  118.  
  119. signed main() {
  120.     ios_base::sync_with_stdio(false);
  121.     cin.tie(0);
  122.    
  123.     // int n;
  124.     // cin >> n;
  125.     // while (true) {
  126.     //  gen(n);
  127.     // }
  128.  
  129.     string s;
  130.     cin >> s;
  131.     cout << solve(s) << endl;
  132. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement