最大化 N 叉树中每个节点的 MEX 值的总和

在本问题中，我们有一个 N 叉树，每个节点都有一个值，其取值范围在 [0, n-1] 之间。定义 MEX 值为最小非负整数中没有出现在该节点的子节点中的最小值。例如，如果一个节点有子节点 0，1，3，则其 MEX 值为 2。

我们的目标是选择每个节点的值，使得所有节点的 MEX 值的总和最大化。在该问题中，我们可以将每个节点的值视为已知，并且寻找每个节点的子节点的值。

解决方案

我们可以使用 DP 解决此问题。我们使用 $f_i(j)$ 表示该节点的值为 $j$，在子树中所有节点的 MEX 值之和。我们可以将该问题分成两部分：

1. 当前节点为 $j$ 时，处理其所有子节点的值；
2. 尝试将子节点的值分配给每个子节点，然后计算其 MEX 值的总和。

考虑第一部分，对于每个子节点 $k$，存在两种情况：

1. 当前节点 $j$ 的值不在子节点 $k$ 的子树中，这种情况不会影响 MEX 值的计算。
2. 当前节点 $j$ 的值在子节点的子树中，这种情况会使得子节点的 MEX 值增加。因此，对于每个子节点，我们选择使得它的 MEX 值增加最大的值来计算。

考虑第二部分，对于每个子节点 $k$，将子节点的值分配给每个子节点，并计算其 MEX 值的总和。为了最大化 MEX 值的总和，我们需要选择使 MEX 增加最大的值。具体而言，我们可以排序子节点，然后顺序地分配子节点的值。

递归方程为：

$$ f_i(j) = \max_{k=0}^{n-1} \left(\frac{d_{j,k}+M}{M+1}\left(f_{i,k}(j) + \sum_{m \in C_k} (f_{i,m}(j)+1)\right)\right) $$

其中，$C_k$ 表示子树 $k$ 中的所有节点。

代码实现

from collections import defaultdict

def mex(node, nodes):
    m = 0
    while m in nodes:
        m += 1
    return m

def dfs(node, parent):
    # 预处理本节点的子节点
    children = defaultdict(list)
    for child in graph[node]:
        if child != parent:
            subtree = dfs(child, node)
            children[subtree].append(child)

    # 预处理子树中的所有节点
    nodes = set()
    for subtree in children:
        nodes |= subtree

    # 计算 MEX 值
    res = 0
    for i in range(len(children)):
        m = mex(i, nodes)
        res += m
        nodes.add(m)

    # DP 状态转移
    dp = [[0] * (n+1) for _ in range(len(children)+1)]
    for i, subtree in enumerate(children):
        for j in range(n):
            for k in range(n):
                if k in subtree:
                    dp[i+1][j] += dp[i][j]
                    continue
                t = mex(j, subtree)
                if t == k:
                    dp[i+1][j] += dp[i][j]
                else:
                    k_res = dp[i][j] + res
                    nodes |= subtree
                    for m in subtree:
                        k_res += (dp[i][m]+1)
                    dp[i+1][t] = max(dp[i+1][t], k_res)
                    nodes -= subtree

    return nodes

n = int(input())
graph = defaultdict(list)
for i in range(n-1):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)

dfs(1, 0)