[Codeforces 715C] Digit Tree

[题目链接]

         https://codeforces.com/contest/715/problem/C

[算法]

        考虑点分治 

        一条路径(x , y)合法当且仅当 : d(x) * 10 ^ dep(x) + d(y) = 0(mod m) , 其中d(u)表示u到分治重心路径上数字拼接起来所形成的数

        统计答案时 ， 我们只需维护一个map ， 维护10 ^ -dep(u) * d(u) (mod m)

        然后计算每个点的贡献即可

        时间复杂度 : O(NlogN ^ 2)

[代码]

       

#include<bits/stdc++.h>
using namespace std;
#define N 100010
typedef long long ll;
typedef long double ld;
typedef unsigned long long ull;

struct edge
{
        int to , w , nxt;
} e[N << 1];

int n , m , tot , len , root;
ll ans;
int pw[N] , head[N] , size[N] , weight[N] , D[N] , depth[N];
bool visited[N];
map<int , int> mp;

template <typename T> inline void chkmax(T &x,T y) { x = max(x,y); }
template <typename T> inline void chkmin(T &x,T y) { x = min(x,y); }
template <typename T> inline void read(T &x)
{
    T f = 1; x = 0;
    char c = getchar();
    for (; !isdigit(c); c = getchar()) if (c == '-') f = -f;
    for (; isdigit(c); c = getchar()) x = (x << 3) + (x << 1) + c - '0';
    x *= f;
}
inline void addedge(int u , int v , int w)
{
        ++tot;
        e[tot] = (edge){v , w , head[u]};
        head[u] = tot;
}
inline void getroot(int u , int par , int total)
{
        size[u] = 1;
        weight[u] = 0;
        for (int i = head[u]; i; i = e[i].nxt)
        {
                int v = e[i].to;
                if (v == par || visited[v]) continue;
                getroot(v , u , total);
                size[u] += size[v];
                chkmax(weight[u] , size[v]);    
        }        
        chkmax(weight[u] , total - size[u]);
        if (weight[u] < weight[root]) root = u;
}
inline void exgcd(int a , int b , int &x , int &y)
{
        if (b == 0)
        {
                x = 1;
                y = 0;
        } else
        {
                exgcd(b , a % b , y , x);
                y -= a / b * x;
        }
}
inline int inv(int a)
{
        int x , y;
        exgcd(a , m , x , y);
        return (x % m + m) % m;
}
inline void dfs(int u , int par , int d1 , int d2)
{
        if (depth[u] > 0) ++mp[(1ll * d2 % m * inv(pw[depth[u]] % m)) % m];
        D[++len] = d1;
        for (int i = head[u]; i; i = e[i].nxt)
        {
                int v = e[i].to , w = e[i].w;
                if (v == par || visited[v]) continue;
                depth[v] = depth[u] + 1;
                dfs(v , u , (1ll * w * pw[depth[v] - 1] % m + d1) % m , (10ll * d2 % m + w) % m);
        }
}
inline ll calc(int u , int d)
{
        mp.clear();
        len = 0;
        if (!d) dfs(u , -1 , 0 , 0);
        else dfs(u , -1 , d % m , d % m);
        ll res = 0;
        for (int i = 1; i <= len; ++i)
        {
                int goal = ((m - D[i]) % m + m) % m;
                res += (ll)mp[goal];
                if (!d && !D[i]) ++res;
        }
        return res;
}
inline void work(int u)
{
        visited[u] = true;
        depth[u] = 0;
        ans += calc(u , 0);
        for (int i = head[u]; i; i = e[i].nxt)
        {
                int v = e[i].to , w = e[i].w;
                if (visited[v]) continue;
                depth[v] = 1;
                ans -= calc(v , w);
        }
        for (int i = head[u]; i; i = e[i].nxt)
        {
                int v = e[i].to;
                if (visited[v]) continue;
                root = 0;
                getroot(v , u , size[v]);
                work(root);
        }
}

int main()
{
        
        read(n); read(m);
        pw[0] = 1;
        for (int i = 1; i <= n; ++i) pw[i] = 1ll * pw[i - 1] * 10 % m;
        for (int i = 1; i < n; ++i)
        {
                int u , v , w;
                read(u); read(v); read(w);
                ++u; ++v;
                addedge(u , v , w);
                addedge(v , u , w);
        }
        weight[0] = n;
        root = 0;
        getroot(1 , 0 , n);
        work(root);
        ans -= n;
        printf("%I64d\n" , ans);
        
        return 0;
    
}