Tsr and Variance

题目描述

Tsr is a cute boy with handsome moustache.

You are given a sequence with length $n$. Tsr wants you to calculate the sum of variance of each successive subsequence. Note: The variance in this problem should’t be divided by length.

Recall $\overline{a}_{l, r}=\frac{1}{r-l+1} \sum_{i=l}^r a_i$. Then you are supposed to calculate $\sum_{i=1}^n \sum_{j=i}^n \sum_{k=i}^j (a_k-\overline{a}_{i, j})^2$.

题意概述

给定一个长度为$n$的序列，求$\sum_{i=1}^n \sum_{j=i}^n \sum_{k=i}^j (a_k-\overline{a}_{i, j})^2$。有$T$组数据。

数据范围：$1 \le T \le 20, \; 1 \le n \le 10000, \; 1 \le a_i \le 10$。

算法分析

$$\begin{align} \sum_{i=1}^n \sum_{j=i}^n \sum_{k=i}^j (a_k-\overline{a}_{i, j})^2 &= \sum_{i=1}^n \sum_{j=i}^n \sum_{k=i}^j (a_k^2-2a_k \overline{a}_{i, j}+\overline{a}_{i, j}^2) \\ &= \sum_{k=1}^n \sum_{i=1}^k \sum_{j=k}^n a_k^2-2\sum_{i=1}^n \sum_{j=i}^n \overline{a}_{i, j} \sum_{k=i}^j a_k+ \sum_{i=1}^n \sum_{j=i}^n (j-i+1) \overline{a}_{i, j}^2 \\ &= \sum_{k=1}^n k(n-k+1)a_k^2-\sum_{i=1}^n \sum_{j=i}^n (j-i+1) \overline{a}_{i, j}^2 \end{align}$$

前一项可以在$O(n)$时间内求出，只需考虑如何求后一项。令$s_i=\sum_{j=1}^i a_j$。

$$\begin{align} \sum_{i=1}^n \sum_{j=i}^n (j-i+1) \overline{a}_{i, j}^2 &= \sum_{j=1}^n \sum_{i=1}^j \frac{(s_j-s_{i-1})^2}{j-i+1} \\ &= \sum_{j=1}^n \sum_{i=1}^j \frac{s_j^2-2s_js_{i-1}+s_{i-1}^2}{j-i+1} \\ &= \sum_{j=1}^n s_j^2 \sum_{i=1}^j \frac{1}{i}-2\sum_{j=1}^n s_j \sum_{i=1}^j \frac{s_{i-1}}{j-(i-1)}+\sum_{j=1}^n \sum_{i=1}^j \frac{s_{i-1}^2}{j-(i-1)} \end{align}$$

第一项也可以在$O(n)$时间内求出，而后面两项都是卷积的形式，可以用FFT在$O(n\log n)$时间内求出。

代码

#include <cmath>
#include <cstdio>
#include <cstring>
#include <algorithm>

int const N = 10005, T = 400005;
double const PI = acos(-1);

int a[N], rev[T];

struct Point {
    double x, y;
    Point(double _x = 0, double _y = 0) : x(_x), y(_y) {}
    friend Point const operator + (Point const &a, Point const &b) {
        return Point(a.x + b.x, a.y + b.y);
    }
    friend Point const operator - (Point const &a, Point const &b) {
        return Point(a.x - b.x, a.y - b.y);
    }
    friend Point const operator * (Point const &a, Point const &b) {
        return Point(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
    }
    friend Point const operator / (Point const &a, double const &n) {
        return Point(a.x / n, a.y / n);
    }
} wn[T], A[T], B[T], C[T];

int init(int n) {
    int m = n, l = 0;
    for (n = 1; n <= m; n <<= 1, ++l) ;
    for (int i = 1; i < n; ++i) {
        rev[i] = rev[i >> 1] >> 1 | (i & 1) << l - 1;
    }
    for (int i = 0; i < n >> 1; ++i) {
        wn[i] = Point(cos(2 * PI / n * i), sin(2 * PI / n * i));
    }
    return n;
}

void fft(Point *a, int n, int inv = 0) {
    for (int i = 0; i < n; ++i) {
        if (i < rev[i]) {
            std::swap(a[i], a[rev[i]]);
        }
    }
    for (int i = 1; i < n; i <<= 1) {
        for (int j = 0; j < n; j += i << 1) {
            for (int k = 0; k < i; ++k) {
                Point x = a[j + k], y = wn[n / (i << 1) * k] * a[j + k + i];
                a[j + k] = x + y;
                a[j + k + i] = x - y;
            }
        }
    }
    if (inv) {
        std::reverse(a + 1, a + n);
        for (int i = 0; i < n; ++i) {
            a[i] = a[i] / n;
        }
    }
}

int main() {
    int T;
    scanf("%d", &T);
    for (; T--;) {
        int n;
        scanf("%d", &n);
        double ans = 0;
        for (int i = 1; i <= n; ++i) {
            scanf("%d", &a[i]);
            ans += 1. * a[i] * a[i] * i * (n - i + 1);
            a[i] += a[i - 1];
        }
        double sum = 0;
        for (int i = 1; i <= n; ++i) {
            sum += 1. / i;
            ans -= sum * a[i] * a[i];
        }
        int len = init(n << 1);
        for (int i = 0; i < n; ++i) {
            A[i] = Point(a[i]);
            B[i] = Point(1. * a[i] * a[i]);
            C[i] = Point(1. / (i + 1));
        }
        for (int i = n; i < len; ++i) {
            A[i] = B[i] = C[i] = Point(0);
        }
        fft(A, len);
        fft(B, len);
        fft(C, len);
        for (int i = 0; i < len; ++i) {
            A[i] = A[i] * C[i];
            B[i] = B[i] * C[i];
        }
        fft(A, len, 1);
        fft(B, len, 1);
        for (int i = 0; i < n; ++i) {
            ans += 2 * a[i + 1] * A[i].x - B[i].x;
        }
        printf("%.10f\n", ans);
    }
    return 0;
}