RT,感觉自己代码的时空复杂度都假完了,但是不卡常以最大点不到 300 ms 的速度过去了,感觉很神奇.

大概思路是建虚树,对于一种颜色,断掉点后会变成很多个连通块,每个连通块的点都要加上总点数减这个连通块点数的贡献.然后如果当前点就是这种颜色的关键点,贡献就是总点数.

然后我的做法是每个连通块的顶端都一定是某个关键点的儿子,所以将其与其儿子加入虚树后,不难发现每个连通块都能通过一段 DFS 序区间挖掉若干个区间后形成,然后用差分加即可.

虽然所有挖掉区间的总个数是 $O(n)$ 的,但是考虑到我虚树上有许多非关键节点,所以暴力合并这些挖掉区间复杂度感觉是假的,自认为的正确做法是对于每个区间找到它上面第一个包含它的区间.

但是我试着写了一法暴力合并,发现很快的就过去了,有大佬能来解释一下这个复杂度到底是不是对的啊!

Code:

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

#define int long long
#define fir first
#define sec second
#define mkp make_pair 
#define pb push_back
#define lep( i, l, r ) for ( int i = ( l ); i <= ( r ); ++ i )
#define rep( i, r, l ) for ( int i = ( r ); i >= ( l ); -- i )
 
typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;
typedef pair < int, int > pii;
 
char _c; bool _f; template < class type > inline void read ( type &x ) {
	_f = 0, x = 0;
	while ( _c = getchar (), !isdigit ( _c ) ) if ( _c == '-' ) _f = 1;
	while ( isdigit ( _c ) ) x = x * 10 + _c - '0', _c = getchar (); if ( _f ) { x = -x; }
}

template < class type > inline void chkmin ( type &x, type y ) { x = ( x <= y ? x : y ); }
template < class type > inline void chkmax ( type &x, type y ) { x = ( x >= y ? x : y ); }

const int N = 1e5 + 5;

int n, m, idx;
int dfn[N], f[N][21], st[N][21], dep[N], a[N], siz[N];
int vis[N];

vector < pair < int, int > > tmp, seg[N];
vector < int > q[N], p[N];

void dfs ( int x, int fa ) {
  dfn[x] = ++ idx, siz[x] = 1;
  dep[x] = dep[fa] + 1;
  f[x][0] = fa;
  for ( int i = 1; i <= 18; i ++ ) {
    f[x][i] = f[f[x][i - 1]][i - 1];
  }
  for ( int v : p[x] ) {
    if ( v != fa ) {
      dfs ( v, x );
      siz[x] += siz[v];
    }
  }
}

int LCA ( int x, int y ) {
  if ( dep[x] > dep[y] ) {
    swap ( x, y );
  }
  for ( int i = 18; i >= 0; i -- ) {
    if ( dep[f[y][i]] >= dep[x] ) {
      y = f[y][i];
    }
  }
  if ( x == y ) {
    return x;
  }
  for ( int i = 18; i >= 0; i -- ) {
    if ( f[x][i] != f[y][i] ) {
      x = f[x][i], y = f[y][i];
    }
  }
  return f[x][0];
}

void Vtree ( vector < pair < int, int > > k ) {
  sort ( k.begin (), k.end () );
  int lim = k.size ();
  for ( int i = 0; i < lim - 1; i ++ ) {
    int elca = LCA ( k[i].second, k[i + 1].second );
    k.push_back ( make_pair ( dfn[elca], elca ) );
  }
  sort ( k.begin (), k.end () );
  vector < pair < int, int > > :: iterator it = unique ( k.begin (), k.end () );
  k.erase ( it, k.end () );
  for ( int i = 0; i < k.size () - 1; i ++ ) {
    int elca = LCA ( k[i].second, k[i + 1].second );
    q[elca].push_back ( k[i + 1].second );
  }
}

int diff[N];
vector < int > c[N];

void merge ( int x, int v ) {
  for ( pii i : seg[v] ) {
    seg[x].push_back ( i );
  }
}

void dfs2 ( int x ) {
  for ( int v : q[x] ) {
    dfs2 ( v );
    merge ( x, v );
  }
  if ( vis[x] == 2 ) {
    sort ( seg[x].begin (), seg[x].end () );
    int lst = dfn[x], ed = dfn[x] + siz[x] - 1, cnt = 0;
    for ( pii i : seg[x] ) {
      cnt += i.first - lst;
      lst = i.second + 1;
    }
    cnt += ed - lst + 1;
    lst = dfn[x];
    for ( pii i : seg[x] ) {
      diff[lst] += n - cnt, diff[i.first] -= n - cnt;
      lst = i.second + 1;
    }
    diff[lst] += n - cnt, diff[ed + 1] -= n - cnt;
  }
  else if ( vis[x] == 1 ) {
    diff[dfn[x]] += n, diff[dfn[x] + 1] -= n;
    seg[x].clear ();
    seg[x].push_back ( { dfn[x], dfn[x] + siz[x] - 1 } );
  }
}

void clear ( int x ) {
  for ( int v : q[x] ) {
    clear ( v );
  }
  q[x].clear (), seg[x].clear (), vis[x] = 0;
}

void Solve () {
  cin >> n;
  for ( int i = 1; i <= n; i ++ ) {
    cin >> a[i];
    c[a[i]].push_back ( i );
  }
  for ( int i = 1; i < n; i ++ ) {
    int u, v;
    cin >> u >> v;
    p[u].push_back ( v ), p[v].push_back ( u );
  }
  dfs ( 1, 0 );
  for ( int i = 1; i < N; i ++ ) {
    if ( !c[i].size () ) {
      continue;
    }
    // int i = 3;
    tmp.clear ();
    for ( int x : c[i] ) {
      tmp.push_back ( { dfn[x], x } );
      vis[x] = 1;
      for ( int v : p[x] ) {
        if ( vis[v] != 1 && v != f[x][0] ) {
          tmp.push_back ( { dfn[v], v } );
          vis[v] = 2;
        }
      }
    }
    tmp.push_back ( { dfn[1], 1 } );
    if ( vis[1] != 1 ) {
      vis[1] = 2;
    }
    Vtree ( tmp );
    dfs2 ( 1 );
    clear ( 1 );
  }
  for ( int i = 1; i <= n; i ++ ) {
    diff[i] += diff[i - 1];
  }
  for ( int i = 1; i <= n; i ++ ) {
    cout << diff[dfn[i]] << '\n';
  }
}

signed main () {
#ifdef judge
  freopen ( "Code.in", "r", stdin );
  freopen ( "Code.out", "w", stdout );
  freopen ( "Code.err", "w", stderr );
#endif
	Solve ();
	return 0;
}
*/