// Witold Milewski // PA 2024 #include <bits/stdc++.h> using namespace std; #define FOR(i, a, b) for(int i=a; i<=b; ++i) #pragma once #define eb emplace_back #define sz(A) (int)(A.size()) #define ll long long ll mov=0; // ########## POCZATEK https://github.com/kth-competitive-programming/kactl/blob/main/content/numerical/FastFourierTransform.h /** * Author: Ludo Pulles, chilli, Simon Lindholm * Date: 2019-01-09 * License: CC0 * Source: http://neerc.ifmo.ru/trains/toulouse/2017/fft2.pdf (do read, it's excellent) Accuracy bound from http://www.daemonology.net/papers/fft.pdf * Description: fft(a) computes $\hat f(k) = \sum_x a[x] \exp(2\pi i \cdot k x / N)$ for all $k$. N must be a power of 2. Useful for convolution: \texttt{conv(a, b) = c}, where $c[x] = \sum a[i]b[x-i]$. For convolution of complex numbers or more than two vectors: FFT, multiply pointwise, divide by n, reverse(start+1, end), FFT back. Rounding is safe if $(\sum a_i^2 + \sum b_i^2)\log_2{N} < 9\cdot10^{14}$ (in practice $10^{16}$; higher for random inputs). Otherwise, use NTT/FFTMod. * Time: O(N \log N) with $N = |A|+|B|$ ($\tilde 1s$ for $N=2^{22}$) * Status: somewhat tested * Details: An in-depth examination of precision for both FFT and FFTMod can be found * here (https://github.com/simonlindholm/fft-precision/blob/master/fft-precision.md) */ typedef vector<ll> vi; typedef complex<long double> C; typedef vector<long double> vd; void fft(vector<C>& a) { int n = sz(a), L = 31 - __builtin_clz(n); static vector<complex<long double>> R(2, 1); static vector<C> rt(2, 1); // (^ 10% faster if double) for (static int k = 2; k < n; k *= 2) { R.resize(n); rt.resize(n); auto x = polar(1.0L, acos(-1.0L) / k); FOR(i,k,2*k-1) rt[i] = R[i] = i&1 ? R[i/2] * x : R[i/2]; } vi rev(n); FOR(i,0,n-1) rev[i] = (rev[i / 2] | (i & 1) << L) / 2; FOR(i,0,n-1) if (i < rev[i]) swap(a[i], a[rev[i]]); for (int k = 1; k < n; k *= 2) for (int i = 0; i < n; i += 2 * k) FOR(j,0,k-1) { // C z = rt[j+k] * a[i+j+k]; // (25% faster if hand-rolled) /// include-line auto x = (long double *)&rt[j+k], y = (long double *)&a[i+j+k]; /// exclude-line C z(x[0]*y[0] - x[1]*y[1], x[0]*y[1] + x[1]*y[0]); /// exclude-line a[i + j + k] = a[i + j] - z; a[i + j] += z; } } void conv(vd& a, const vd& b) { if (a.empty() || b.empty()) return; vd res(sz(a) + sz(b) - 1); int L = 32 - __builtin_clz(sz(res)), n = 1 << L; vector<C> in(n), out(n); copy(a.begin(), a.end(), begin(in)); FOR(i,0,sz(b)-1) in[i].imag(b[i]); fft(in); for (C& x : in) x *= x; FOR(i,0,n-1) out[i] = in[-i & (n - 1)] - conj(in[i]); fft(out); FOR(i,0,sz(res)-1) res[i] = imag(out[i]) / (4 * n); a=res; } // ########## KONIEC https://github.com/kth-competitive-programming/kactl/blob/main/content/numerical/FastFourierTransform.h const int maxn=502; ll n, A[maxn], pref[maxn]; int32_t main() { cin.tie(0) -> ios_base::sync_with_stdio(0); cin >> n; FOR(i, 1, n) cin >> A[i]; ll maxv = -1e9; ll minv = 1e9; ll ile_0 = 0; FOR(i, 1, n) { pref[i]=pref[i-1]+A[i]; FOR(j, 1, i) { minv=min(minv, pref[i]-pref[j-1]); maxv=max(maxv, pref[i]-pref[j-1]); if(pref[i]-pref[j-1]==0) ++ile_0; } } mov=-minv; if(minv<0) maxv+=mov;; vector<ll> wspol(maxv+1, 0); FOR(i, 1, n) FOR(j, 1, i) ++wspol[pref[i]-pref[j-1]+mov]; if(minv>=0) { // dodatnie ll res = ile_0*ile_0*ile_0; ll cnt = ile_0 + (ile_0*(ile_0-1))*3; cout << (res-cnt)/6 << '\n'; } else { // ujemne vd odp(wspol.begin(), wspol.end()); vd wspo(wspol.begin(), wspol.end()); if(sz(odp)>=3*mov+1) odp.resize(3*mov+1); conv(odp, wspo); conv(odp, wspo); ll res = round(odp[3*mov]); ll cnt=0; if(mov<=sz(wspol)-1) cnt=wspol[mov]; FOR(i, 0, sz(wspol)-1) { ll b = (3*mov)-2*i; if(b<0 || b>sz(wspol)-1) continue; if(i==mov && wspol[i]>=1) cnt += (wspol[i]*(wspol[i]-1))*3; else cnt += wspol[i]*wspol[b]*3; } cout << (res-cnt)/6 << '\n'; } }
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 111 112 113 114 115 116 117 118 119 | // Witold Milewski // PA 2024 #include <bits/stdc++.h> using namespace std; #define FOR(i, a, b) for(int i=a; i<=b; ++i) #pragma once #define eb emplace_back #define sz(A) (int)(A.size()) #define ll long long ll mov=0; // ########## POCZATEK https://github.com/kth-competitive-programming/kactl/blob/main/content/numerical/FastFourierTransform.h /** * Author: Ludo Pulles, chilli, Simon Lindholm * Date: 2019-01-09 * License: CC0 * Source: http://neerc.ifmo.ru/trains/toulouse/2017/fft2.pdf (do read, it's excellent) Accuracy bound from http://www.daemonology.net/papers/fft.pdf * Description: fft(a) computes $\hat f(k) = \sum_x a[x] \exp(2\pi i \cdot k x / N)$ for all $k$. N must be a power of 2. Useful for convolution: \texttt{conv(a, b) = c}, where $c[x] = \sum a[i]b[x-i]$. For convolution of complex numbers or more than two vectors: FFT, multiply pointwise, divide by n, reverse(start+1, end), FFT back. Rounding is safe if $(\sum a_i^2 + \sum b_i^2)\log_2{N} < 9\cdot10^{14}$ (in practice $10^{16}$; higher for random inputs). Otherwise, use NTT/FFTMod. * Time: O(N \log N) with $N = |A|+|B|$ ($\tilde 1s$ for $N=2^{22}$) * Status: somewhat tested * Details: An in-depth examination of precision for both FFT and FFTMod can be found * here (https://github.com/simonlindholm/fft-precision/blob/master/fft-precision.md) */ typedef vector<ll> vi; typedef complex<long double> C; typedef vector<long double> vd; void fft(vector<C>& a) { int n = sz(a), L = 31 - __builtin_clz(n); static vector<complex<long double>> R(2, 1); static vector<C> rt(2, 1); // (^ 10% faster if double) for (static int k = 2; k < n; k *= 2) { R.resize(n); rt.resize(n); auto x = polar(1.0L, acos(-1.0L) / k); FOR(i,k,2*k-1) rt[i] = R[i] = i&1 ? R[i/2] * x : R[i/2]; } vi rev(n); FOR(i,0,n-1) rev[i] = (rev[i / 2] | (i & 1) << L) / 2; FOR(i,0,n-1) if (i < rev[i]) swap(a[i], a[rev[i]]); for (int k = 1; k < n; k *= 2) for (int i = 0; i < n; i += 2 * k) FOR(j,0,k-1) { // C z = rt[j+k] * a[i+j+k]; // (25% faster if hand-rolled) /// include-line auto x = (long double *)&rt[j+k], y = (long double *)&a[i+j+k]; /// exclude-line C z(x[0]*y[0] - x[1]*y[1], x[0]*y[1] + x[1]*y[0]); /// exclude-line a[i + j + k] = a[i + j] - z; a[i + j] += z; } } void conv(vd& a, const vd& b) { if (a.empty() || b.empty()) return; vd res(sz(a) + sz(b) - 1); int L = 32 - __builtin_clz(sz(res)), n = 1 << L; vector<C> in(n), out(n); copy(a.begin(), a.end(), begin(in)); FOR(i,0,sz(b)-1) in[i].imag(b[i]); fft(in); for (C& x : in) x *= x; FOR(i,0,n-1) out[i] = in[-i & (n - 1)] - conj(in[i]); fft(out); FOR(i,0,sz(res)-1) res[i] = imag(out[i]) / (4 * n); a=res; } // ########## KONIEC https://github.com/kth-competitive-programming/kactl/blob/main/content/numerical/FastFourierTransform.h const int maxn=502; ll n, A[maxn], pref[maxn]; int32_t main() { cin.tie(0) -> ios_base::sync_with_stdio(0); cin >> n; FOR(i, 1, n) cin >> A[i]; ll maxv = -1e9; ll minv = 1e9; ll ile_0 = 0; FOR(i, 1, n) { pref[i]=pref[i-1]+A[i]; FOR(j, 1, i) { minv=min(minv, pref[i]-pref[j-1]); maxv=max(maxv, pref[i]-pref[j-1]); if(pref[i]-pref[j-1]==0) ++ile_0; } } mov=-minv; if(minv<0) maxv+=mov;; vector<ll> wspol(maxv+1, 0); FOR(i, 1, n) FOR(j, 1, i) ++wspol[pref[i]-pref[j-1]+mov]; if(minv>=0) { // dodatnie ll res = ile_0*ile_0*ile_0; ll cnt = ile_0 + (ile_0*(ile_0-1))*3; cout << (res-cnt)/6 << '\n'; } else { // ujemne vd odp(wspol.begin(), wspol.end()); vd wspo(wspol.begin(), wspol.end()); if(sz(odp)>=3*mov+1) odp.resize(3*mov+1); conv(odp, wspo); conv(odp, wspo); ll res = round(odp[3*mov]); ll cnt=0; if(mov<=sz(wspol)-1) cnt=wspol[mov]; FOR(i, 0, sz(wspol)-1) { ll b = (3*mov)-2*i; if(b<0 || b>sz(wspol)-1) continue; if(i==mov && wspol[i]>=1) cnt += (wspol[i]*(wspol[i]-1))*3; else cnt += wspol[i]*wspol[b]*3; } cout << (res-cnt)/6 << '\n'; } } |