#include <iostream> #include <vector> #include <algorithm> #include <complex> #include <map> using namespace std; using cd = complex<double>; const double PI = acos(-1); void fft(vector<cd> & a, bool invert) { int n = a.size(); if (n == 1) return; vector<cd> a0(n / 2), a1(n / 2); for (int i = 0; 2 * i < n; i++) { a0[i] = a[2*i]; a1[i] = a[2*i+1]; } fft(a0, invert); fft(a1, invert); double ang = 2 * PI / n * (invert ? -1 : 1); cd w(1), wn(cos(ang), sin(ang)); for (int i = 0; 2 * i < n; i++) { a[i] = a0[i] + w * a1[i]; a[i + n/2] = a0[i] - w * a1[i]; if (invert) { a[i] /= 2; a[i + n/2] /= 2; } w *= wn; } } int32_t main(){ ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0); int n; cin >> n; vector<int> ciag(n); for(int i = 0; i < n; i++){ cin >> ciag[i]; } int min = 1000000000; int max = -1000000000; map<int, int> mapa; for (int i = 0; i < n; i++){ int suma = 0; for (int j = i; j < n; j++){ suma += ciag[j]; if (suma < min){ min = suma; } if (suma > max){ max = suma; } mapa[suma] += 1; } } if (min > 0){ cout << 0 << "\n"; return 0; }else if (max < 0){ cout << 0 << "\n"; return 0; } vector<cd> wspolczynniki(max - min + 1, 0); for (auto x : mapa){ wspolczynniki[x.first - min] = x.second; } //dopelnienie wspolczynnikow int pom = wspolczynniki.size(); for (int i = 0; i < 3*pom; i++){ wspolczynniki.push_back(0); } //dopelnienie wspolczynnikow do potegi 2 while((wspolczynniki.size() & (wspolczynniki.size() - 1)) != 0){ wspolczynniki.push_back(0); } fft(wspolczynniki, false); for (int i = 0; i < wspolczynniki.size(); i++){ wspolczynniki[i] = wspolczynniki[i] * wspolczynniki[i] * wspolczynniki[i]; } fft(wspolczynniki, true); //odczytanie wspolczynniku przy -3*min long long wynik = round(wspolczynniki[-3*min].real()); //cout << wynik << "\n"; wynik -= mapa[0]*mapa[0]*mapa[0]; wynik += mapa[0]*(mapa[0]-1)*(mapa[0]-2); for (auto p : mapa){ if (p.first != 0){ int q = -2*p.first; wynik -= 3*p.second*p.second*mapa[q]; wynik += 3*p.second*(p.second-1)*mapa[q]; } } cout << wynik/6 << "\n"; return 0; }
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | #include <iostream> #include <vector> #include <algorithm> #include <complex> #include <map> using namespace std; using cd = complex<double>; const double PI = acos(-1); void fft(vector<cd> & a, bool invert) { int n = a.size(); if (n == 1) return; vector<cd> a0(n / 2), a1(n / 2); for (int i = 0; 2 * i < n; i++) { a0[i] = a[2*i]; a1[i] = a[2*i+1]; } fft(a0, invert); fft(a1, invert); double ang = 2 * PI / n * (invert ? -1 : 1); cd w(1), wn(cos(ang), sin(ang)); for (int i = 0; 2 * i < n; i++) { a[i] = a0[i] + w * a1[i]; a[i + n/2] = a0[i] - w * a1[i]; if (invert) { a[i] /= 2; a[i + n/2] /= 2; } w *= wn; } } int32_t main(){ ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0); int n; cin >> n; vector<int> ciag(n); for(int i = 0; i < n; i++){ cin >> ciag[i]; } int min = 1000000000; int max = -1000000000; map<int, int> mapa; for (int i = 0; i < n; i++){ int suma = 0; for (int j = i; j < n; j++){ suma += ciag[j]; if (suma < min){ min = suma; } if (suma > max){ max = suma; } mapa[suma] += 1; } } if (min > 0){ cout << 0 << "\n"; return 0; }else if (max < 0){ cout << 0 << "\n"; return 0; } vector<cd> wspolczynniki(max - min + 1, 0); for (auto x : mapa){ wspolczynniki[x.first - min] = x.second; } //dopelnienie wspolczynnikow int pom = wspolczynniki.size(); for (int i = 0; i < 3*pom; i++){ wspolczynniki.push_back(0); } //dopelnienie wspolczynnikow do potegi 2 while((wspolczynniki.size() & (wspolczynniki.size() - 1)) != 0){ wspolczynniki.push_back(0); } fft(wspolczynniki, false); for (int i = 0; i < wspolczynniki.size(); i++){ wspolczynniki[i] = wspolczynniki[i] * wspolczynniki[i] * wspolczynniki[i]; } fft(wspolczynniki, true); //odczytanie wspolczynniku przy -3*min long long wynik = round(wspolczynniki[-3*min].real()); //cout << wynik << "\n"; wynik -= mapa[0]*mapa[0]*mapa[0]; wynik += mapa[0]*(mapa[0]-1)*(mapa[0]-2); for (auto p : mapa){ if (p.first != 0){ int q = -2*p.first; wynik -= 3*p.second*p.second*mapa[q]; wynik += 3*p.second*(p.second-1)*mapa[q]; } } cout << wynik/6 << "\n"; return 0; } |