Gym101002E:K-Inversions

\(Gym101002E:K-Inversions\)

题意描述：

• 题目连接链接

• 给定一个长度为\(N\)只包含\(AB\)的字符串，某个\(A\)的位置为\(j\)，某个\(B\)的位置为\(i\)，求\(j-k=k\)的数对有多少个，输出\(k=1,2,...,n-1\)的情况。

数据范围：

• \(1\leq n\leq 10^6\)。

思路：

• \(FFT\)。
• 假设\(A\)的位置为\(x\)，\(B\)的位置为\(y\)，则题目需要\(x-y=k\)。
• 如果说我们将\(A,B\)的位置看作是多项式的幂，就可以在\(nlogn\)的时间内求出所有幂为\(x+y\)的系数。
• 所以这里需要转换一下，令\(y=(n-y-1)\)，那么\(x-y=x-n+y+1=k\)，即\(x+y=n+k-1\)。
• 所以将\(B\)的位置反转，套用\(fft\)求解。

#include<bits/stdc++.h>
using namespace std;
const int maxn = 6e6 + 10;
const double PI = acos(-1.0);
char s[maxn];
int limit, l, r[maxn];
struct Complex
{
    double x, y;
    Complex(double xx = 0, double yy = 0){
        x = xx, y = yy;
    }
    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);
    }
}a[maxn], b[maxn];

void fft(Complex c[], int type)
{
    for(int i = 0; i < limit; i++)
        if(i < r[i]) swap(c[i], c[r[i]]);
    for(int mid = 1; mid < limit; mid <<= 1)
    {
        Complex wn(cos(PI/mid), type*sin(PI/mid));
        for(int R = mid<<1, j = 0; j < limit; j+= R)
        {
            Complex w(1, 0);
            for(int k = 0; k < mid; k++, w = w*wn)
            {
                Complex x = c[j+k], y = w*c[j+mid+k];
                c[j+k] = x+y;
                c[j+mid+k] = x - y;
            }
        }
    }
}

int main()
{
    scanf("%s", s);
    int n = strlen(s);
    for(int i = 0; i < n; i++)
    {
        if(s[i] == 'A') a[i].x = 1;
        else b[n-i-1].x = 1;
    } limit = 1;
    while(limit <= n+n) limit <<= 1, l++;
    for(int i = 0; i < limit; i++)
        r[i] = (r[i>>1]>>1)|((i&1)<<(l-1));
    fft(a, 1), fft(b, 1);
    for(int i = 0; i <= limit; i++)
        a[i] = a[i]*b[i];
    fft(a, -1);
    for(int i = n; i < n+n-1; i++)
        printf("%d\n", (int)(a[i].x/limit+0.5));

    return 0;
}