codeforces716E （点分治）

Problem Digit Tree

题目大意

　　给一棵树，有边权1~9。

　　询问有多少个点对（i，j），将i--j路径上的数字依次连接后所形成新数字可以被k整除。gcd（K,10）=1

解题分析

　　点分治。考虑某一次分治，根为rt，求出所有子节点到根所形成数字为A，根到所有子节点所形成数字为B。

　　那么即求出所有满足 ( A[i] * 10 ^ B[j] . len + B[j] ) % K == 0的点对。

　　转化后为 A[i] == (K - B[j]) * 10 ^ ( - B[j] . len ) ， 用map处理一下即可。

参考程序

#include <map>
#include <set>
#include <stack>
#include <queue>
#include <cmath>
#include <ctime>
#include <string>
#include <vector>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cassert>
#include <iostream>
#include <algorithm>
#pragma comment(linker,"/STACK:102400000,102400000")
using namespace std;

#define N 100008    
#define LL long long
#define lson l,m,rt<<1    
#define rson m+1,r,rt<<1|1 
#define clr(x,v) memset(x,v,sizeof(x));
#define bitcnt(x) __builtin_popcount(x)
#define rep(x,y,z) for (int x=y;x<=z;x++)
#define repd(x,y,z) for (int x=y;x>=z;x--)
const int mo  = 1000000007;
const int inf = 0x3f3f3f3f;
const int INF = 2000000000;
/**************************************************************************/ 

int n,m,phi;
int lt[N],size[N],f[N],vis[N],a[N],b[N],sum,tot,root,pre1[N],pre2[N];
LL ans,Alltmp;
map <int,int> mp;
struct edge{int u,v,w,nt;}eg[N*2];
void add(int u,int v,int w){
    eg[++sum]=(edge){u,v,w,lt[u]}; lt[u]=sum;
}
LL quick(LL x,LL y)
{
    LL res=1;
    while (y)
    {
        if  (y & 1) res=res*x % m;
        x=x*x % m;
        y>>=1;
    }
    return res;
}
void init()
{
    int mm=m;
    phi=m;
    for (int i=2;i*i<=m;i++)
    {
        if (mm % i==0)
        {
            while (mm % i==0) mm/=i;
            phi=phi/i*(i-1);
        }
    }
    if (mm>1) phi=phi/mm*(mm-1);
    clr(lt,0); sum=1; 
    clr(f,0); f[0]=INF;
    clr(vis,0);
    ans=0;
    pre1[0]=1; pre2[0]=1;
    rep(i,1,100000)
    {
        pre1[i]=1ll*pre1[i-1]*10 % m;
        pre2[i]=quick(pre1[i],phi-1);
    }
}
void getRoot(int u,int fa)
{    
    size[u]=1; f[u]=0;
    for (int i=lt[u];i;i=eg[i].nt)
    {
        int v=eg[i].v;
        if (vis[v] || v==fa) continue;
        getRoot(v,u);
        size[u]+=size[v];
        f[u]=max(f[u],size[v]);
    }
    f[u]=max(f[u],tot-size[u]);
    if (f[u]<f[root]) root=u;
}
void getA(int u,int fa,int len)
{    
    for (int i=lt[u];i;i=eg[i].nt)
    {
        int v=eg[i].v;
        if (vis[v] || v==fa) continue;
        a[v]=(1ll*eg[i].w*pre1[len]+a[u]) % m;
        getA(v,u,len+1);
    }
    Alltmp+=mp[a[u]];
    int tp=(m-b[u]) % m;
    tp=1ll*tp*pre2[len] % m;
    if (a[u]==tp) Alltmp--;
}
void getB(int u,int fa,int len)
{
    int tp=(m-b[u]) % m;
    tp=1ll*tp*pre2[len]% m;
    mp[tp]++;
    for (int i=lt[u];i;i=eg[i].nt)
    {
        int v=eg[i].v;
        if (vis[v] || v==fa) continue;
        b[v]=(1ll*b[u]*10+eg[i].w) % m;
        getB(v,u,len+1);
    }
}
LL calc(int u,int key)
{
    key%=m;
    Alltmp=0; mp.clear();
    a[u]=b[u]=key;
    getB(u,0,key==0?0:1); getA(u,0,key==0?0:1); 
    return Alltmp;
}
void solve(int u)
{
    ans+=calc(u,0); vis[u]=1;
    for (int i=lt[u];i;i=eg[i].nt)
    {
        int v=eg[i].v;
        if (vis[v]) continue; 
        ans-=calc(v,eg[i].w);
        root=0; tot=size[v];
        getRoot(v,u);
        solve(root);
    }
}
int main()
{
    scanf("%d%d",&n,&m);
    init();
    rep(i,1,n-1)
    {
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        add(u+1,v+1,w);
        add(v+1,u+1,w);
    }
    root=0; tot=n;
    getRoot(1,0);
    solve(root);
    cout<<ans<<endl;
}