【题解】HackerRank - Recurrent on a Tree

Posted on Edited on In Disqus:

题目

https://www.hackerrank.com/contests/w34/challenges/recurrent-on-tree

题意

给一棵树，每个点有一个值。求任意两个顶点（可以相同）路径上的值的和经过 f 函数后的和。f(x) 是斐波那契数列第 x 项。

题解

  • 斐波那契数列的某一项可以通过矩阵快速幂求得。与之类似的，利用矩阵乘法可以将 f(? + ?) 的问题转化为求 Mat(?) * Mat(?) 的问题，进而转化为寻常的树上 dp。
  • 由于运算中矩阵都是某一个矩阵的幂，所以满足交换律。
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
#include <bits/stdc++.h>
#define FOR(i, x, y) for (decay<decltype(y)>::type i = (x), _##i = (y); i < _##i; ++i)
#define FORD(i, x, y) for (decay<decltype(x)>::type i = (x), _##i = (y); i > _##i; --i)
using namespace std;
typedef long long LL;
const LL MOD = 1E9 + 7;
const int maxm = 1E5 + 1;
const int maxn = 2E5 + 10;
struct Mat {
    LL v[2][2];
    inline void clear() {
        v[0][0] = v[0][1] = v[1][0] = v[1][1] = 0;
    }
};
void init();
Mat operator * (Mat x, Mat y);
Mat operator + (Mat x, Mat y);
Mat operator - (Mat x, Mat y);

Mat r[maxm], s[maxn], ans, cs[maxn];
vector<int> G[maxn];
int n, u, v, a[maxn];

void dfs(int k, int fa) {
    FOR(i, 0, G[k].size())
        if (G[k][i] != fa)
            dfs(G[k][i], k);
    FOR(i, 0, G[k].size()) {
        int t = G[k][i];
        if (t == fa) continue;
        ans = ans + r[a[k]] * s[t] * (cs[k] - s[t]);
        s[k] = s[k] + s[t];
    }
    s[k] = s[k] * r[a[k]] + r[a[k]];
    cs[fa] = cs[fa] + s[k];
    ans = ans + s[k] + s[k] - r[a[k]];
}

int main() {
    init();
    cin >> n;
    FOR(i, 1, n) {
        scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    FOR(i, 1, n + 1) scanf("%d", &a[i]);
    dfs(1, 0);
    cout << (ans.v[1][0] + ans.v[1][1]) % MOD << endl;
}

void init() {
    Mat p;
    p.v[0][0] = 1; p.v[0][1] = 1;
    p.v[1][0] = 1; p.v[1][1] = 0;
    r[0].v[0][0] = r[0].v[1][1] = 1;
    FOR(i, 1, maxm)
        r[i] = r[i - 1] * p;
}

Mat operator * (Mat x, Mat y) {
    Mat ret;
    FOR(i, 0, 2)
        FOR(j, 0, 2) {
            ret.v[i][j] = 0;
            FOR(k, 0, 2)
                ret.v[i][j] = (ret.v[i][j] + x.v[i][k] * y.v[k][j]) % MOD;
        }
    return ret;
}

Mat operator + (Mat x, Mat y) {
    Mat ret;
    FOR(i, 0, 2)
        FOR(j, 0, 2)
            ret.v[i][j] = (x.v[i][j] + y.v[i][j]) % MOD;
    return ret;
}

Mat operator - (Mat x, Mat y) {
    Mat ret;
    FOR(i, 0, 2)
        FOR(j, 0, 2)
            ret.v[i][j] = (x.v[i][j] - y.v[i][j] + MOD) % MOD;
    return ret;
}