题解「猫猫虫上树」

· · 个人记录

人话翻译:对于某一深度找到两条链,使其对应深度元素之积的和最大。

学过线性代数的同学会发现,这本质上就是两个向量的内积。我们定义树的链向量为从根到底所有元素形成的有序组:\mathbf{x_{dep}}=\begin{pmatrix} \mathbf{x_{dep_1}}\ \mathbf{x_{dep_2}}\ \cdots \ \mathbf{x_{dep_n}}\ \end{pmatrix}

内积的定义同理:\begin{bmatrix}\mathbf{x},\mathbf{y}\end{bmatrix}=\mathbf{x}_1\mathbf{y}_1+\mathbf{x}_2\mathbf{y}_2+\cdots+\mathbf{x}_n\mathbf{y}_n

那么我们只需要对一个固定的向量 \mathbf{x},找到一个存在的 \mathbf{y},使其内积最大。显然他只要与自己内积,因为 \begin{bmatrix}\mathbf{x},\mathbf{y}\end{bmatrix} \leq \begin{bmatrix}\mathbf{x},\mathbf{x}\end{bmatrix}

所以我们只需要记录当前深度链向量模长的最大值即可,时间复杂度 O(n+q)

const int mod = 998244353;
const int N = 2e6 + 5;

int a[N], fa[N], dep[N]; ll mx[N];
int n, q, x, y; vector <int> g[N];

inline void dfs(int u, int father, ll val) {
    dep[u] = dep[father] + 1;
    val += 1ll * a[u] * a[u];
    chkmax(mx[dep[u]], val);
    for (int v : g[u]) {
        if (v == father) continue;
        dfs(v, u, val);
    }
}

signed main(void) {
    read(n), read(q);
    for (int i = 1; i <= n; i++) read(a[i]);
    for (int i = 2; i <= n; i++) read(fa[i]), g[fa[i]].push_back(i);
    dfs(1, 0, 0);
    while (q--) {
        read(x);
        writeln(mx[x] % mod);
    }
    return 0;
}