【题解】[GDKOI2021] 空中恋爱论

气象万千

思路

cdq 分治 + FFT。

首先意识到对和满足要求的区间计数，可以先求前缀和，再转化成前缀和的端点对计数。

令 \(A_n = \sum\limits_{i = 1}^n a_i, B_n = \sum\limits_{i = 1}^n b_i\).

那么对于一个点对 \(i, j (i \leq j)\)，如果 \(B_j - B_i > 0\)，就说明 \((i, j]\) 是合法的区间。

现在的问题转化成统计 \(B\) 中的每一个点对。

这个问题看起来没有其他做法，考虑分治。

假设当前区间为 \([l, r]\). 我们需要统计的点对满足 \(0 \leq l \leq i \leq \lfloor \frac{l + r}{2} \rfloor < j \leq r \leq n\)，并且 \(B_j - B_i > 0\)。在此基础上，它会对 \(A_j - A_i\) 的位置贡献 \(1\).

那么先把数组按照 \(B\) 排序，保证分治时右半段区间的 \(B\) 之和大于等于左半段。然后把 \([l, mid]\) 中的点值设为 \(-A_i\)，\((mid, r]\) 中的点值设为 \(A_j\). 只需要做一个无标号计数就行。随便 FFT 或者 NTT 带走。

考虑到每次都做 FFT 会浪费很多时间，考虑类似根号分治，当区间长度大于阈值时才 FFT。

阈值取 \(\sqrt{n \log n}\) 最优，时间复杂度 \(O(n \sqrt{n \log n})\).

代码

// FFT 板子掉精，我觉得该反思人品了
#include <cstdio>
#include <cmath>
#include <complex>
#include <iostream>
#include <algorithm>
using namespace std;
 
#define fi first
#define se second
 
typedef long long ll;
typedef double db;
 
const int maxn = 1e6 + 5;
const ll inf = 1e18;
const db PI = acos(-1);
 
int n;
int rev[maxn];
ll ans[maxn];
pair<ll, ll> a[maxn];
complex<double> F[maxn], G[maxn];
 
void calc_rev(int k) { for (int i = 0; i < k; i++) rev[i] = (rev[i >> 1] >> 1 | (i & 1 ? k >> 1 : 0)); }
 
// void FFT(complex<double> *A, int n)
// {
//     for (int i = 1; i < n; i++)
//         if (rev[i] > i) swap(A[i], A[rev[i]]);
//     for (int len = 2, m = 1; len <= n; m = len, len <<= 1)
//     {
//         complex<double> W(cos(PI / m), sin(PI / m)), w(1.0, 0.0);
//         for (int l = 0, r = len - 1; r <= n; l += len, r += len)
//         {
//             auto w0 = w;
//             for (int p = l; p < l + m; p++)
//             {
//                 auto x = A[p] + w0 * A[p + m], y = A[p] - w0 * A[p + m];
//                 A[p] = x, A[p + m] = y;
//                 w0 *= W;
//             }
//         }
//     }
// }
 
void FFT(complex<double>* c, int n){
    calc_rev(n);
    for (int i = 0; i < n; i++) if (i < rev[i]) swap(c[i], c[rev[i]]);
    for (int i = 1;i < n; i <<= 1)
    {
        complex<double> omn(cos(PI / i), sin(PI / i));
        for (int j = 0; j < n; j += (i << 1))
        {
            complex<double> om(1, 0);
            for (int k = 0; k < i; k++, om = om * omn)
            {
                complex<double> x = c[j + k], y = c[j + k + i] * om;
                c[j + k] = x + y, c[j + k + i] = x - y;
            }
        }
    }
}
 
void IFFT(complex<double> *A, int n)
{
    FFT(A, n), reverse(A + 1, A + n);
    for (int i = 0; i < n; i++) A[i] /= n;
}
 
void cdq(int l, int r)
{
    if (l >= r) return;
    int mid = (l + r) >> 1;
    ll lmn = inf, lmx = -inf, rmn = inf, rmx = -inf;
    cdq(l, mid), cdq(mid + 1, r);
    if (r - l <= 1e4)
    {
        for (int i = l; i <= mid; i++)
            for (int j = mid + 1; j <= r; j++)
                ans[a[j].se - a[i].se + n]++;
        return;
    }
    for (int i = l; i <= mid; i++) lmn = min(lmn, a[i].se), lmx = max(lmx, a[i].se);
    for (int i = mid + 1; i <= r; i++) rmn = min(rmn, a[i].se), rmx = max(rmx, a[i].se);
    int k = 1;
    while (k <= lmx - lmn + rmx - rmn) k <<= 1;
    for (int i = 0; i <= k; i++) F[i] = (0, 0);
    for (int i = l; i <= mid; i++)
    {
        double tmp = F[lmx - a[i].se].real() + 1.0;
        F[lmx - a[i].se].real(tmp);
    }
    for (int i = mid + 1; i <= r; i++)
    {
        double tmp = F[a[i].se - rmn].imag() + 1.0;
        F[a[i].se - rmn].imag(tmp);
    }
    FFT(F, k);
    for (int i = 0; i < k; i++) F[i] *= F[i];
    IFFT(F, k);
    for (int i = 0; i <= lmx - lmn + rmx - rmn; i++) ans[i + rmn - lmx + n] += (int)(F[i].imag() / 2.0 + 0.5);
}
 
signed main()
{
    scanf("%lld", &n);
    for (int i = 1; i <= n; i++) scanf("%lld", &a[i].se), a[i].se += a[i - 1].se;
    for (int i = 1; i <= n; i++) scanf("%lld", &a[i].fi), a[i].fi += a[i - 1].fi;
    sort(a, a + n + 1);
    cdq(0, n);
    for (int i = 0; i <= (n << 1); i++) printf("%lld ", ans[i]);
    return 0;
}