CF 993E Nikita and Order Statistics

传送门

题目大意

给你一个数组 a1∼n a 1 ∼ n $a_{1 \sim n}$，对于 k=0∼n k = 0 ∼ n $k = 0 \sim n$，求出有多少个数组上的区间满足：区间内恰好有 k k $k$ 个数比 x$x$$x$ 小。 x x $x$ 为一个给定的数。

n≤2×105$n\le 2\times {10}^5$$n \le 2 \times 10^5$。值域没有意义。

思路

一眼看到标签写着 FFT……

把所有小于 x x $x$ 的数记为 1$1$$1$，剩下的数记为 0 0 $0$，则题目要求的是有多少个区间有 k$k$$k$ 个 1 1 $1$。对这个 01$01$$01$ 数组求前缀和，设为 s s $s$，则题目要求的是有多少个 (l,r)(l≤r)$(l,r)(l\le r)$$(l, r) \pod {l \le r}$ 满足 sr−sl−1=k s r − s l − 1 = k $s_r - s_{l - 1} = k$。

现在相当于有一个长度为 n+1​ n + 1 ​ $n + 1​$ 的数组 s​ s ​ $s​$，满足 si≤si+1≤si+1​ s i ≤ s i + 1 ≤ s i + 1 ​ $s_i \le s_{i + 1} \le s_i + 1​$，问两个元素的差的取值的情况数。设多项式 f​ f ​ $f​$ 的 xi​ x i ​ $x^i​$ 项的系数等于 sk=i​ s k = i ​ $s_k = i​$ 的 k​ k ​ $k​$ 的个数，设多项式 g​ g ​ $g​$ 的 x−i​ x − i ​ $x^{-i}​$ 项的系数等于 sk=i​ s k = i ​ $s_k = i​$ 的 k​ k ​ $k​$ 的个数。对 f​ f ​ $f​$ 和 g​ g ​ $g​$ 做卷积，则积的 xk(k>0)​ x k ( k > 0 ) ​ $x^k \pod {k > 0}​$ 的系数就是对应的答案。当 k=0​ k = 0 ​ $k = 0​$ 时，首先要减去自己卷上自己的 n+1​ n + 1 ​ $n + 1​$ 次，然后再除以二，就是对应的答案。

时间复杂度 O(nlogn) O ( n log ⁡ n ) $O(n \log n)$。

参考代码

#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <cstring>
#include <cassert>
#include <cctype>
#include <climits>
#include <ctime>
#include <iostream>
#include <algorithm>
#include <vector>
#include <string>
#include <stack>
#include <queue>
#include <deque>
#include <map>
#include <set>
#include <bitset>
#include <list>
#include <functional>
typedef long long LL;
typedef unsigned long long ULL;
using std::cin;
using std::cout;
using std::endl;
typedef LL INT_PUT;
INT_PUT readIn()
{
    INT_PUT a = 0; bool positive = true;
    char ch = getchar();
    while (!(ch == '-' || std::isdigit(ch))) ch = getchar();
    if (ch == '-') { positive = false; ch = getchar(); }
    while (std::isdigit(ch)) { a = a * 10 - (ch - '0'); ch = getchar(); }
    return positive ? -a : a;
}
void printOut(INT_PUT x)
{
    char buffer[20]; int length = 0;
    if (x < 0) putchar('-'); else x = -x;
    do buffer[length++] = -(x % 10) + '0'; while (x /= 10);
    do putchar(buffer[--length]); while (length);
}

typedef long double REAL;
const REAL PI = std::acos((REAL)-1);
const int maxn = int(2e5) + 5;
int n, x;
int a[maxn];

struct Complex
{
    REAL x, y;
    Complex() {}
    Complex(REAL x, REAL y) : x(x), y(y) {}
    Complex operator+(const Complex& b) const
    {
        return Complex(x + b.x, y + b.y);
    }
    Complex operator-(const Complex& b) const
    {
        return Complex(x - b.x, y - b.y);
    }
    Complex operator*(const Complex& b) const
    {
        return Complex(x * b.x - y * b.y, x * b.y + y * b.x);
    }
    void operator/= (const REAL& b)
    {
        x /= b;
        y /= b;
    }
};
void DFT(Complex* a, int logn, bool IDFT)
{
    int pre = -1;
    static int revbit[1 << 20];
    int n = 1 << logn;
    if (logn != pre)
    {
        pre = logn;
        revbit[0] = 0;
        for (int i = 1; i < n; i++)
            revbit[i] = (revbit[i >> 1] >> 1) | ((i & 1) << (logn - 1));
    }
    for (int i = 0; i < n; i++)
        if (i < revbit[i])
            std::swap(a[i], a[revbit[i]]);

    for (int i = 1; i <= logn; i++)
    {
        int S = 1 << i;
        int half = S >> 1;
        Complex w1(std::cos(2 * PI / S), std::sin(2 * PI / S) * (IDFT ? -1 : 1));
        for (int j = 0; j < n; j += S)
        {
            Complex* A = a + j;
            Complex w(1, 0);
            for (int k = 0; k < half; k++)
            {
                Complex t = A[k + half] * w;
                A[k + half] = A[k] - t;
                A[k] = A[k] + t;
                w = w * w1;
            }
        }
    }

    if (IDFT)
        for (int i = 0; i < n; i++)
            a[i] /= n;
}

Complex A[1 << 20];
Complex B[1 << 20];

void run()
{
    n = readIn();
    x = readIn();
    for (int i = 1; i <= n; i++)
        a[i] = a[i - 1] + (readIn() < x);

    for (int i = 0; i <= n; i++)
    {
        A[n + a[i]].x++;
        B[n - a[i]].x++;
    }
    int logn = 0;
    while (1 << logn < n * 3) logn++;
    DFT(A, logn, false);
    DFT(B, logn, false);
    int N = 1 << logn;
    for (int i = 0; i < N; i++)
        A[i] = A[i] * B[i];
    DFT(A, logn, true);
    LL t;
    t = A[n << 1].x + 0.5;
    printOut((t - n - 1) >> 1);
    for (int i = 1; i <= n; i++)
    {
        putchar(' ');
        printOut(A[(n << 1) + i].x + 0.5);
    }
}

int main()
{
    run();
    return 0;
}

要是标签都不看就会做就好了。唉，谁叫我太弱了呢 QAQ。

一开始居然把 t 定义成 int 类型了，妥妥的爆零啊 QAQ。