找到要从树中删除的边以最大化组件的 XOR 乘积

介绍

给定一棵有 $n$ 个节点的无根树，每个节点 $i$ 上有一个权值 $a_i$。你需要找到一条边 $(u, v)$ 并删除它，使得树被分成两个子树，每个子树上节点的权值的异或积最大。返回这个最大的异或积。

需要注意的是，如果节点 $u$ 和 $v$ 之间有多条边，只需要删除其中一个即可。

算法思路

这里提供一种基于 DFS 的思路。

对于一条从根节点到某个叶子节点的路径上的所有节点，这些节点的路径异或总和就是这些节点的权值异或积。因此可以从根节点开始 DFS，对于每个节点 $u$，计算出 $u$ 到子树内所有节点的权值异或积并记录下来。对于一条边 $(u, v)$，假设 $u$ 是 $v$ 的父亲，可以得到两个子树的权值异或积：$u$ 的子树权值异或积和 $v$ 子树权值异或积。那么删除边 $(u, v)$ 后，这两个子树的权值异或积就分别为 $u$ 到 $v$ 路径上的节点的权值异或积和 $v$ 到根节点路径上的节点的权值异或积。因此可以在 DFS 的过程中同时记录下这两个值，找到最大的乘积即可。

具体地，对于一个节点 $u$，假设 $v_1, v_2, \dots, v_k$ 是 $u$ 的子节点，$\operatorname{XOR}(i,j)$ 表示 $i$ 和 $j$ 的异或和。那么 $u$ 到子树内所有节点的权值异或积可以递归计算得到：

$$ w_u = a_u \operatorname{XOR}{i=1}^k w{v_i} $$

同时假设 $w_{u,1}$ 表示 $u$ 到 $v_1$ 路径上的节点的权值异或积，$w_{u,2}$ 表示 $v_1$ 到根节点路径上的节点的权值异或积，那么可以得到：

$$ \begin{aligned} w_{u,1} &= a_u \operatorname{XOR}{i=2}^k w{v_i} \ w_{u,2} &= w_{v_1,2} \operatorname{XOR} (w_v \operatorname{XOR} w_{u,1}) \operatorname{XOR} (a_u \operatorname{XOR} w_{u,1}) \ \end{aligned} $$

其中 $v$ 表示节点 $v_1$ 的父亲。注意到可以对每个节点 $u$ 计算 $w_{u,1}$ 和 $w_{u,2}$，并记录下来，最后可以根据这两个值计算出删除每条边后的异或积并返回最大值。

代码实现

以下代码实现了上述思路。

class Solution:
    def solve(self, n: int, edges: List[List[int]], w: List[int]) -> int:
        g = [[] for _ in range(n + 1)]
        for u, v in edges:
            g[u].append(v)
            g[v].append(u)
        vis = [False] * (n + 1)
        ans = 0

        def dfs(u: int, p: int) -> Tuple[int, int]:  # 返回 (w1, w2)
            vis[u] = True
            wu1, wu2 = w[u], 0
            for v in g[u]:
                if v != p:
                    wv1, wv2 = dfs(v, u)
                    wu1 ^= wv1
                    if wv2 > wu2:
                        wu1, wu2 = wv2 ^ wu1 ^ w[u], wv2
            ans = max(ans, wu1 * wu2)
            return wu1, wu2 ^ w[u]

        for u in range(1, n + 1):
            if not vis[u]:
                dfs(u, 0)
        return ans

其中 dfs 函数的参数 u 表示当前节点，p 表示父节点。在计算子树权值异或积的时候，需要将结果初始化为 w[u] 并递归计算。在计算 $w_{u,2}$ 的时候需要考虑不选 $v_1$ 的情况，因此需要先计算出去掉 $v_1$ 后的 $w_{u,1}$ 并存储在 wu1 变量中。