题解:P14468 [COCI 2025/2026 #1] 和谐 / Harmonija

· · 题解

思路

考虑到路径上每次选择红蓝颜色时,红蓝节点的数量差受到强限制 \vert{}x - y\vert{} \le 2,状态空间极小且转移具有结合律,容易想到利用 (\max, +) 矩阵乘法结合树上倍增来维护动态规划状态的转移。

设当前路径前缀中红点数量为 x,蓝点数量为 y。题目要求在任意时刻均满足 \max(x, y) < \min(x, y) + 3,这等价于 \vert{}x - y\vert{} \le 2。因此 x - y 的合法取值仅有 -2, -1, 0, 1, 2 这 5 种。我们将 x - y + 2 作为状态下标,分别映射到 0, 1, 2, 3, 4。

对于节点 u 的染色选择:

  1. 染成红色:x 增加 1,x - y 增加 1,状态从 j 转移到 j + 1,获得权值 c_u。
  2. 染成蓝色:y 增加 1,x - y 减少 1,状态从 j 转移到 j - 1,获得权值 p_u。

在 (\max, +) 广义矩阵乘法下,可构造 5 \times 5 的基础转移矩阵 W_u,其中转移权值定义为 W_{u, j+1, j} = c_u 与 W_{u, j, j+1} = p_u(0 \le j \le 3),其余位置初始化为 -\infty。

定义状态向量为 1 \times 5 的行向量,代表 5 种不同的 x - y 状态值(分别对应状态 -2, -1, 0, 1, 2):

\begin{bmatrix} f_0 & f_1 & f_2 & f_3 & f_4 \end{bmatrix} \times \begin{bmatrix} -\infty & c_u & -\infty & -\infty & -\infty \\ p_u & -\infty & c_u & -\infty & -\infty \\ -\infty & p_u & -\infty & c_u & -\infty \\ -\infty & -\infty & p_u & -\infty & c_u \\ -\infty & -\infty & -\infty & p_u & -\infty \end{bmatrix} = \begin{bmatrix} f'_0 & f'_1 & f'_2 & f'_3 & f'_4 \end{bmatrix}

由于矩阵乘法不满足交换律,自底向上与自顶向下的转移顺序不同,需要分别维护:

处理询问 (u, v) 时,求出 l = \text{LCA}(u, v)。初始时 x = 0, y = 0(对应状态 2),权值为 0。根据 l 的位置分类讨论:

与初始矩阵相乘后,取所有合法终态的最大值即为答案。

时间复杂度为 O((n + q) \cdot 5^3 \log n),空间复杂度为 O(n \cdot 5^2 \log n)。

Code

#include <bits/stdc++.h>
#define int long long
using namespace std;

namespace Creeper{
    const int N = 1e5 + 5, inf = -1e18;
    int n, q, c[N], p[N], fa[20][N], dep[N];
    vector<int> G[N];
    struct mtx {
        int a[5][5];
        void clear() {
            for (int i = 0; i <= 4; i++) for (int j = 0; j < 5; j++) a[i][j] = inf;     
        }
        void print() {
            for (int i = 0; i <= 4; i++) {
                for (int j = 0; j <= 4; j++) {
                    cout << a[i][j] << ' '; 
                }
                cout << '\n';
            }
            cout << '\n';
        }
    }w[N], f[18][N], g[18][N];
    mtx operator*(const mtx a, const mtx b)  {
        mtx res;
        res.clear();
        for (int i = 0; i <= 4; i++) {
            for (int j = 0; j <= 4; j++) {
                for (int k = 0; k <= 4; k++) {
                    res.a[i][j] = max(res.a[i][j], a.a[i][k] + b.a[k][j]);
                }
            }
        }
        return res;
    }
    void dfs(int u, int pre) {
        fa[0][u] = pre;
        dep[u] = dep[pre] + 1;
        f[0][u] = w[pre];
        g[0][u] = w[pre];
        for (auto v: G[u]) {
            if (v == pre) continue;
            dfs(v, u);
        }
    }
    int lca(int u, int v) {
        if (dep[u] < dep[v]) swap(u, v);
        for (int i = 17; i >= 0; i--) {
            if (dep[fa[i][u]] >= dep[v]) {
                u = fa[i][u];
            }
        }
        if (u == v) return u;
        for (int i = 17; i >= 0; i--) {
            if (fa[i][u] != fa[i][v]) {
                u = fa[i][u];
                v = fa[i][v];
            }
        }
        return fa[0][u];
    }
    mtx kth(int u, int k, int o) {
        mtx res;
        res.clear();
        if (o == 0) {
            res = w[u];
            for (int i = 17; i >= 0; i--) {
                if ((k >> i) & 1) {
                    res = res * f[i][u];
                    u = fa[i][u];
                }
            }
            return res;
        } else {
            res = w[u];
            for (int i = 17; i >= 0; i--) {
                if ((k >> i) & 1) {
                    res = g[i][u] * res;
                    u = fa[i][u];
                }
            }
            return res;
        }
    }
    void solve() {
        cin >> n >> q;
        for (int i = 1; i <= n; i++) cin >> c[i];
        for (int i = 1; i <= n; i++) cin >> p[i];
        w[0].clear();
        for (int i = 1; i <= n; i++) {
            w[i].clear();
            for (int j = 0; j <= 3; j++) w[i].a[j + 1][j] = c[i];
            for (int j = 0; j <= 3; j++) w[i].a[j][j + 1] = p[i];
        }
        for (int i = 1; i < n; i++) {
            int u, v;
            cin >> u >> v;
            G[u].push_back(v);
            G[v].push_back(u); 
        }
        dfs(1, 0);
        for (int i = 1; i <= 17; i++) {
            for (int j = 1; j <= n; j++) {
                fa[i][j] = fa[i - 1][fa[i - 1][j]];
                f[i][j] = f[i - 1][j] * f[i - 1][fa[i - 1][j]];
                g[i][j] = g[i - 1][fa[i - 1][j]] * g[i - 1][j];
            }
        }
        while (q--) {
            int u, v;
            cin >> u >> v;
            int l = lca(u, v);
            mtx ans;
            ans.clear();
            ans.a[1][2] = 0;
            if (l == u) {
                mtx x = kth(v, dep[v] - dep[l], 1);
                ans = ans * x;
            } else if (l == v) {
                mtx x = kth(u, dep[u] - dep[l], 0);
                ans = ans * x;              
            } else {
                mtx x = kth(u, dep[u] - dep[l], 0);
                mtx y = kth(v, dep[v] - dep[l] - 1, 1); 
                ans = ans * x;
                ans = ans * y; 
            }
            cout << max({ans.a[1][0], ans.a[1][1], ans.a[1][2], ans.a[1][3], ans.a[1][4]}) << '\n';
        }
    }
} 

signed main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    int T = 1;
//  cin >> T;
    while (T--) Creeper::solve();
}