终于好好写了一次点分题！

3.14 模拟赛 C

Tree

考虑一个题目描述的三角形其实就是同色三角形的个数，是个经典问题，可以看成总方案数减去异色三角形数，于是问题变成了统计异色三角形的数量。我们知道一个异色三角形必然有两个边同色，我们考虑一个三角形 $u \to v , v\to t , u \to t$ （原题是有向的），那么我们在 $u\to v,u\to t$ 的位置统计下 $u$ ，再在 $v\to t,u\to t$ 的位置统计下 $t$ ，显然这两次统计一次统计同色一次统计异色，于是相当于只需要考虑 $u\to v$ 优秀的 $v$ 的数量和 $v\to t$ 优秀的 $t$ 的数量，于是我们就有了个 $O(n^2)$ 的优秀算法。

但是可以统计的过程考虑点分，考虑对 $u$ 统计共有多少个 $v$ 满足 $u \to v$ 优秀。

设当前点分中心 $r$ ，我们对于一个子树内的点 $u$ 要找之前子树有多少个 $v$ 满足

$$f[u] = w[r] + \sum_{x : r\to u} w_xk^{dep[x]}\\g[u] = \sum_{x : r\to u} w_xk^{n-dep[x]}\\f[v] \times k^{dep[u]+1} + \frac{g[u]}{k^{n-dep[u]-1}} = 0$$

减过去，除过来，就是

$$f[v] = -\frac{g[u]}{k^{n}}$$

于是开个 hash 表统计一下就好了。

还得考虑有多少个 $v$ 满足 $v\to u$ 优秀

$$-k^nf[u] = g[v]$$

就是把 $g[v]$ 放进 hash 表。

所以这个东西就可以 $O(n\log n)$ 做辣。。

然后考试时少了个 %P 调了一个小时。。。。

题目很卡常，并且学到了一个卡常技巧，当你访问 map 的一个位置的时候会建立这个位置，所以应当先拿 if(mp.count(a)) 来判一下再用这个位置的值。。

#include "iostream"
#include "algorithm"
#include "cstring"
#include "cstdio"
#include "set"
#include "vector"
#include "map"
#define MAXN 100006
//#define int long long
using namespace std;
typedef long long ll;
const double eps = 1e-7;
int n , k , P;
int w[MAXN] , pw[MAXN] , ivp[MAXN];
vector<int> G[MAXN];
int siz[MAXN] , mx[MAXN] , vis[MAXN] , rt , pr;
inline int Pow( int x , int a ) {
    int cur = x % P , ans = 1;
    while( a ) {
        if( a & 1 ) ans = 1ll * ans * cur % P;
        cur = 1ll * cur * cur % P , a >>= 1;
    }
    return ans;
}
void dfs1( int u , int fa ) {
    siz[u] = 1 , mx[u] = 0;
    for( int v : G[u] ) if( v != fa && !vis[v] ) {
            dfs1( v , u ) , siz[u] += siz[v] , mx[u] = max( mx[u] , siz[v] );
        }
}
void dfs2( int u , int fa ) {
    mx[u] = max( mx[u] , siz[rt] - siz[u] );
    if( mx[u] < mx[pr] ) pr = u;
    for( int v : G[u] ) if( v != fa && !vis[v] )
            dfs2( v , u );
}
int findrt( int u ) {
    pr = rt = u; dfs1( u , u ); dfs2( u , u );
    return pr;
}
int f[MAXN] , g[MAXN] , dep[MAXN];
int to[MAXN] , co[MAXN];
map<int,int> ff , gg;
void dfs3( int u , int fa ) {
    if( ff.count( ( P - 1ll * g[u] * ivp[n] % P ) % P ) )to[u] += ff[( P - 1ll * g[u] * ivp[n] % P ) % P];
    if( gg.count( ( P - 1ll * f[u] * pw[n] % P ) % P ) ) co[u] += gg[( P - 1ll * f[u] * pw[n] % P ) % P];
    for( int v : G[u] ) if( v != fa && !vis[v] ) {
            dep[v] = dep[u] + 1;
            f[v] = ( f[u] + 1ll * w[v] * pw[dep[v]] % P ) % P;
            g[v] = ( g[u] + 1ll * w[v] * pw[n - dep[v]] % P ) % P;
            dfs3( v , u );
        }
}
void dfs4( int u , int fa ) {
    ++ ff[f[u]] , ++ gg[g[u]];
    for( int v : G[u] ) if( v != fa && !vis[v] ) dfs4( v , u );
}
void work( int u ) {
    vis[u] = 1;
    ff.clear() , gg.clear();
    f[u] = w[u]; ++ ff[f[u]] , ++ gg[0] , g[u] = 0;
    for( int v : G[u] ) if( !vis[v] ) {
            dep[v] = 1 , f[v] = ( 1ll * w[v] * k % P + w[u] ) % P , g[v] = 1ll * w[v] * pw[n - 1] % P;
            dfs3( v , u ) , dfs4( v , u );
        }
    to[u] += ff[0];
    co[u] += gg[( P - 1ll * f[u] * pw[n] % P ) % P];
    ff.clear() , gg.clear();
    for( int i = G[u].size() - 1 ; i >= 0 ; -- i ) {
        int v = G[u][i];
        if( vis[v] ) continue;
        dfs3( v , u ) , dfs4( v , u );
    }
    for( int v : G[u] ) if( !vis[v] ) work( findrt( v ) );
}
int main() {
//    freopen("c1.in","r",stdin);
//    freopen("ot","w",stdout);
    cin >> n >> k >> P; k %= P;
    pw[0] = ivp[0] = 1 , pw[1] = k , ivp[1] = Pow( k , P - 2 );
    for( int i = 2 ; i <= n ; ++ i ) pw[i] = 1ll * pw[i - 1] * k % P , ivp[i] = 1ll * ivp[i - 1] * ivp[1] % P;
    for( int i = 1 ; i <= n ; ++ i ) scanf("%d",&w[i]) , w[i] = ( w[i] % P + P ) % P;
    for( int i = 1 , u , v ; i < n ; ++ i ) {
        scanf("%d%d",&u,&v);
        G[u].push_back( v ) , G[v].push_back( u );
    }
    work( findrt( 1 ) );
    long long res = 0;
    for( int i = 1 ; i <= n ; ++ i ) {
//        cout << to[i] << ' ' << co[i] << endl;
        res += 2ll * to[i] * ( n - to[i] ) + 2ll * co[i] * ( n - co[i] );
        res += 1ll * to[i] * ( n - co[i] ) + 1ll * co[i] * ( n - to[i] );
    }
    cout << ( 1ll * n * n * n - res / 2 ) << endl;
}