OJ - Formula Computing

Posted on Edited on In , Views: 35 Disqus:

算法课的 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<105), calculate C[k]=∑i=kn−1(a[i]×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)-th line contains a[i],b[i].

Output

n lines. The k-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

很简单，二重循环直接算即可，复杂度 Θ(n2)：

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]=∑i=kn−1a[i]×b[i−k]=∑j=0n−1−ka′[(n−1−k)−j]×b[j]。至此已经有卷积的雏形了，再做一步代换C′[k]=C[n−1−k]=∑j=0ka′[k−j]×b[j]，终于出现了卷积形式。

有了卷积形式之后就可以使用 FFT 来加速了，复杂度 Θ(nlog⁡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;
}