CSES 2112：One Bit Positions

CSES 2112：One Bit Positions

題目大意：給一個 $0、1$ 組成的字串，輸出每個距離為 $i$ 的兩個 $1$ 有幾對。

解法：$\text{FFT}$。

$\text{Code：}$

#pragma GCC optimize("O3")
#include <bits/stdc++.h>
using namespace std;

using CPX = complex<double>;

const double pi = acos (-1);

int n, siz;
vector<CPX> w, a, b;

void out (vector<CPX> v) {
    for (auto x : v)
        cout << x << ' ';
    cout << '\n';
}

void init() {
    string str;
    cin >> str;
    n = str.size();
    siz = 1 << (__lg (n) + 2);
    w.resize (siz);
    a.resize (siz);
    b.resize (siz);
    for (int i = 0; i < siz; i++) {
        w[i] = {cos (2 * pi * i / siz), sin (2 * pi*i / siz) };
        a[i] = (i >= n ? 0 : str[i] - '0');
        b[i] = (n - i - 1 < 0 ? 0 : str[n - i - 1] - '0');
    }
}

vector<CPX> fft (vector<CPX> v, int ty) {
    int rev[siz] = {};
    for (int i = 0; i < siz; i++) {
        rev[i] = rev[i >> 1] >> 1;
        if (i & 1)
            rev[i] |= siz >> 1;
    }
    for (int i = 0; i < siz; i++) {
        if (i < rev[i])
            swap (v[i], v[rev[i]]);
    }
    for (int len = 2; len <= siz; len <<= 1) {
        int half = len >> 1;
        for (int i = 0; i < siz; i += len) {
            for (int j = 0; j < half; j++) {
                CPX t = v[i + half + j] * CPX (w[siz / len * j].real(), ty * w[siz / len * j].imag());
                v[i + half + j] = v[i + j] - t;
                v[i + j] = v[i + j] + t;
            }
        }
    }
    if (ty == -1) {
        for (int i = 0; i < siz; i++)
            v[i] /= siz;
    }
    return v;
}

void solve() {
    init();
    a = fft (a, 1);
    b = fft (b, 1);
    for (int i = 0; i < siz; i++)
        a[i] *= b[i];
    a = fft (a, -1);
    for (int i = n - 2; i >= 0; i--)
        cout << fixed << setprecision (0) << a[i].real() + 1e-5 << ' ';
    cout << '\n';
}

int main() {
    ios::sync_with_stdio (false), cin.tie (0);
    solve();
}

我的分享就到這裡結束了，如果喜歡我的 $\text{Blog}$，歡迎追蹤！