OJ - Formula Computing

Posted on | Edited on | In , | Comments:

算法课的 project 中有一道挺有意思的题，原文如下：

English Version

Formula Computing

Description

Given two arrays $a$ and $b$ of length $n$ which contains positive integers less than or equal to 100 (we assume that the array index starts at zero and $n < 10^5$), calculate $C[k]=\sum_{i=k}^{n-1}(a[i]\times b[i - k])$ for $k = 0,1,…,n-1$.

Input

The first line contains a positive integer $n$. There are two positive integers in each of the next $n$ lines representing the elements in the two arrays, i.e., the $(i+2)\text{-th}$ line contains $a[i]$,$b[i]$.

Output

$n$ lines. The $k\text{-th}$ line contains one positive number which represents $C[k-1]$.

Sample

  • Input

    1
    2
    3
    4
    5
    6
    5
    3 1
    2 4
    1 1
    2 4
    1 4

  • Output

    1
    2
    3
    4
    5
    24
    12
    10
    6
    1

Solution

Naive Method

很简单，二重循环直接算即可，复杂度 $\Theta (n^2)$：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
#include <iostream>
#include <vector>
#include <stdint.h>

using Number = int64_t;

int main() {
    int n;
    std::vector<Number> a, b, c;

    std::ios::sync_with_stdio(false);

    std::cin >> n;
    a.resize(n);
    b.resize(n);
    c.resize(n, 0);

    for (int i = 0; i < n; ++i) {
        std::cin >> a[i] >> b[i];
    }

    for (int i = 0; i < n; ++i) {
        for (int j = i; j < n; ++j) {
            c[j - i] += a[j] * b[i];
        }
    }

    for (int i = 0; i < n; ++i) {
        std::cout << c[i] << std::endl;
    }

    return 0;
}

FFT Method

其实上述算法已经可以通过所有数据点了（可能是数据比较弱），不过显然不是课程期待的算法。

注意到$C[k]$的计算方式有点像卷积操作，令$a’[k]=a[n - 1 - k]$，则有$C[k]=\sum_{i=k}^{n-1}a[i]\times b[i-k]=\sum_{j=0}^{n-1-k}a’[(n-1-k)-j]\times b[j]$。至此已经有卷积的雏形了，再做一步代换$C’[k]=C[n-1-k]=\sum_{j=0}^{k}a’[k-j]\times b[j]$，终于出现了卷积形式。

有了卷积形式之后就可以使用 FFT 来加速了，复杂度 $\Theta (n \log n)$：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
#include <complex>
#include <vector>
#include <iostream>
#include <math.h>

struct FftContext {

    using Complex = std::complex<double>;

    FftContext(int len)
        : len_(len) {
            if (len_ > 1) {
                int hflen = len_ / 2;
                const double PI = acos(-1);
                double ua = PI / hflen;
                w_.assign(hflen, Complex{});
                for (int i = 0; i < w_.size(); ++i) {
                    double arg = ua * i;
                    w_[i] = Complex{cos(arg), sin(-arg)};
                }
            }
        }

    ~FftContext() { }

    void Forward(std::vector<Complex> &data) const {
        if (len_ == 1) { return; }
        BitRevese(data);
        for (int l = 2; l <= len_; l *= 2) {
            FftLayer(data, l, 0);
        }
    }

    void Backward(std::vector<Complex> &data) const {
        if (len_ == 1) { return; }
        BitRevese(data);
        for (int l = 2; l <= len_; l *= 2) {
            FftLayer(data, l, 1);
        }
        for (int i = 0; i < len_; ++i) {
            data[i] /= len_;
        }
    }

private:
    void BitRevese(std::vector<Complex> &data) const {
        int hflen = len_ / 2;
        for (int i = 1, j = hflen; i < len_ - 1; ++i) {
            if (i < j) {
                std::swap(data[i], data[j]);
            }
            int k = hflen;
            while (j >= k) {
                j -= k;
                k /= 2;
            }
            if (j < k) {
                j += k;
            }
        }
    }

    void FftLayer(std::vector<Complex> &data, int gsize, int rev) const {
        int gcnt = len_ / gsize;
        int hfgs = gsize / 2;
        for (int i = 0; i < len_; i += gsize) {
            for (int j = 0; j < hfgs; ++j) {
                const Complex &wj = w_[j * gcnt];
                Complex wmb = data[i + j + hfgs] * (rev ? std::conj(wj): wj);

                data[i + j + hfgs] = data[i + j] - wmb;
                data[i + j] += wmb;
            }
        }
    }

    int len_;
    std::vector<Complex> w_;
};

uint32_t next_power_of_two(uint32_t x) {
    uint32_t r = 1;
    while (r < x) {
        r *= 2;
    }
    return r;
}

int main() {
    int n, pn;

    std::ios::sync_with_stdio(false);

    std::cin >> n;
    pn = next_power_of_two(n); /* padding to power of two */

    FftContext ctx{pn * 2};
    std::vector<std::complex<double>> a,b;
    a.assign(pn * 2, std::complex<double>{});
    b.assign(pn * 2, std::complex<double>{});

    for (int i = 0; i < n; ++i) {
        int ai, bi;
        std::cin >> ai >> bi;
        a[n - 1 - i] = ai;     /* here a[k] is a'[k] actually */
        b[i] = bi;
    }

    ctx.Forward(a);
    ctx.Forward(b);

    for (int i = 0; i < a.size(); ++i) {
        a[i] *= b[i];
    }

    ctx.Backward(a);           /* here a[k] is c'[k] actually */
    for (int i = 0; i < n; ++i) {
        std::cout << (int)std::round(std::real(a[n - 1 - i])) << std::endl;
    }

    return 0;
}