Solution -「51nod 1868」彩色树

\(\mathcal{Description}\)

  Link & 双倍经验 Link.

  给定一棵 \(n\) 个结点的树，每个结点有一种颜色。记 \(g(u,v)\) 表示 \(u\) 到 \(v\) 简单路径上的颜色种数，求

\[\sum_{\{p_n\}}\sum_{i=1}^{n-1}g(p_i,p_{i+1})
\]

  其中 \(\{p_n\}\) 表示 \(1\sim n\) 的排列。

  \(n\le10^5\)，答案对 \((10^9+7)\) 取模。

\(\mathcal{Solution}\)

  常见但不熟悉的 trick 题。（

  不难想到分颜色计算贡献。而“路径上至少出现某种颜色”不好计算，考虑反向计算“路径上不包含某种颜色”的路径条数。

  对于颜色 \(c\)，删除所有颜色为 \(c\) 的结点，记得到联通块的大小为 \(s_{1..m}\)，那么上述路径数量为 \(\sum_{i=1}^m\frac{s_i(s_i-1)}2\)。称一个联通块深度最小的结点为其顶点，我们尝试在顶点处对这个联通块计数。可以看出，要不顶点是根，要不顶点父亲的颜色为 \(c\)。前者单独考虑，后者仅需用顶点子树大小减去顶点子树内以 \(c\) 颜色的结点作为根的子树大小。所以用前后作差的 trick：DFS 进入子树前，记录目前以 \(c\) 色点为根的子树大小和 \(s\)，退出子树后，得到当前以 \(c\) 色点为根的子树大小和 \(t\)。那么该联通块的大小即为 \(t-s\)。

  所以 DFS 一遍就可以啦。复杂度 \(\mathcal O(n)\)。

\(\mathcal{Code}\)

/* Clearink */

#include <cstdio>
#include <vector>
#include <algorithm>

#define rep( i, l, r ) for ( int i = l, repEnd##i = r; i <= repEnd##i; ++i )
#define per( i, r, l ) for ( int i = r, repEnd##i = l; i >= repEnd##i; --i )

inline int rint () {
	int x = 0, f = 1; char s = getchar ();
	for ( ; s < '0' || '9' < s; s = getchar () ) f = s == '-' ? -f : f;
	for ( ; '0' <= s && s <= '9'; s = getchar () ) x = x * 10 + ( s ^ '0' );
	return x * f;
}

const int MAXN = 1e5, MOD = 1e9 + 7;
int n, ecnt, ans, clr[MAXN + 5], head[MAXN + 5], siz[MAXN + 5], sum[MAXN + 5];
struct Edge { int to, nxt; } graph[MAXN * 2 + 5];

inline int mul ( const long long a, const int b ) { return a * b % MOD; }
inline int sub ( int a, const int b ) { return ( a -= b ) < 0 ? a + MOD : a; }
inline int add ( int a, const int b ) { return ( a += b ) < MOD ? a : a - MOD; }

inline void link ( const int s, const int t ) {
	graph[++ecnt] = { t, head[s] };
	head[s] = ecnt;
}

inline int count ( const int s ) { return ( s * ( s - 1ll ) >> 1 ) % MOD; }

inline void dfs ( const int u, const int fa ) {
	siz[u] = 1, ++sum[clr[u]];
	for ( int i = head[u], v; i; i = graph[i].nxt ) {
		if ( ( v = graph[i].to ) ^ fa ) {
			int s = sum[clr[u]];
			dfs ( v, u ), siz[u] += siz[v];
			int t = sum[clr[u]], dlt = t - s;
			ans = add ( ans, count ( siz[v] - dlt ) );
			sum[clr[u]] += siz[v] - dlt;
		}
	}
}

int main () {
	n = rint ();
	rep ( i, 1, n ) clr[i] = rint ();
	for ( int i = 1, u, v; i < n; ++i ) {
		u = rint (), v = rint ();
		link ( u, v ), link ( v, u );
	}
	dfs ( 1, 0 );
	rep ( c, 1, n ) {
		if ( c ^ clr[1] ) {
			ans = add ( ans, count ( n - sum[c] ) );
		}
	}
	ans = sub ( mul ( n, count ( n ) ), ans );
	int fct = 1;
	rep ( i, 1, n - 2 ) fct = mul ( fct, i );
	printf ( "%d", mul ( ans, mul ( 2, mul ( n - 1, fct ) ) ) );
	return 0;
}