CodeForces 715C Digit Tree (樹的分治)

分析:

求所有滿足情況的鏈的個數,這個時候很容易想到樹的分治,每次求經過rt的所有滿足條件的鏈的個數。由於從rt開始的鏈可以進行合併。所有直接暴力所有子樹就可以了。複雜度就是o(nlogn)

代碼:

#include <bits/stdc++.h>

#define LL long long
#define FOR(i,x,y)  for(int i = x;i < y;i ++)
#define IFOR(i,x,y) for(int i = x;i > y;i --)
#define pb  push_back
#define mp  make_pair

using namespace std;

void gcd(LL a,LL b,LL &d,LL &x,LL &y){
    if(!b)  {d = a;x = 1;y = 0;}
    else    {gcd(b,a%b,d,y,x);y -= x*(a/b);}
}

LL Inv(LL a,LL n){
    LL d,x,y;
    gcd(a,n,d,x,y);
    return d == 1 ? (x+n)%n : -1;
}

const int maxn = 100100;
const int inf = 1<<30;
LL bit[maxn],inv[maxn];

int head[maxn],edge_cnt;

struct  Edge{
    int u,v,nt;
    int w;
}edge[maxn<<1];

void addedge(int u,int v,int w){
    edge[edge_cnt].u = u;
    edge[edge_cnt].v = v;
    edge[edge_cnt].w = w;
    edge[edge_cnt].nt = head[u];
    head[u] = edge_cnt++;
}

bool vis[maxn]; //表明該點有沒有被找過
int sz[maxn],ss[maxn];  //子樹節點總數,子結點的子數中的最大值
int lt[maxn],lt_cnt;    //該子樹內所有子樹標號
int n,m;
LL k;

LL d[maxn];
int d_cnt;  //所有子節點到根的距離

LL ans;

//找重心
void dfs_lt(int u,int fa){//遍歷整個子樹
    lt[lt_cnt++] = u;
    sz[u] = 1;
    ss[u] = 0;
    for(int i = head[u];i != -1;i = edge[i].nt){
        int v = edge[i].v;
        if(v == fa || vis[v])   continue;
        dfs_lt(v,u);
        sz[u] += sz[v];
        ss[u] = max(ss[u],sz[v]);
    }
}

int getroot(int u){
    lt_cnt = 0;
    dfs_lt(u,-1);
    int m_sz = inf,rt;
    FOR(i,0,lt_cnt){
        int tem = max(lt_cnt-sz[lt[i]],ss[lt[i]]);
        if(tem < m_sz)  {m_sz = tem;rt = lt[i];}
    }
    return rt;
}

map <LL,int> L,R;

//計算所有結點到根的距離
LL dfs_11(int u,int fa,int dis,LL w){
    LL p = ((m-w*inv[dis]%m)%m+m)%m;
    LL res = L[p];
    if(!w)  res += 1;
    for(int i = head[u];i != -1;i = edge[i].nt){
        int v = edge[i].v;
        if(v == fa || vis[v]) continue;
        LL tt = (w*10+edge[i].w)%m;
        res += dfs_11(v,u,dis+1,tt);
    }
    return res;
}

void dfs_12(int u,int fa,int dis,LL w){
    ++ L[w%m];
    for(int i = head[u];i != -1;i = edge[i].nt){
        int v= edge[i].v;
        if(v == fa || vis[v]) continue;
        LL tt = (w+edge[i].w*bit[dis]%m)%m;
        dfs_12(v,u,dis+1,tt);
    }
}

LL dfs_21(int u,int fa,int dis,LL w){
    LL p = ((m-w*inv[dis]%m)%m+m)%m;
    LL res = R[p];
    for(int i = head[u];i != -1;i = edge[i].nt){
        int v = edge[i].v;
        if(v == fa || vis[v]) continue;
        LL tt = (w*10+edge[i].w)%m;
        res += dfs_21(v,u,dis+1,tt);
    }
    return res;
}

void dfs_22(int u,int fa,int dis,LL w){
    ++ R[w%m];
    for(int i = head[u];i != -1;i = edge[i].nt){
        int v= edge[i].v;
        if(v == fa || vis[v]) continue;
        LL tt = (w+edge[i].w*bit[dis]%m)%m;
        dfs_22(v,u,dis+1,tt);
    }
}


vector <int> node;

//計算點對
LL calc(int u){
    node.clear();
    LL res = 0;
    L.clear();  R.clear();
    for(int i = head[u];i != -1;i = edge[i].nt){
        int v = edge[i].v;
        node.pb(i);
        if(vis[v])  continue;
        res += dfs_11(v,u,1,edge[i].w%m);
        dfs_12(v,u,1,edge[i].w%m);
    }
    res += L[0];
    reverse(node.begin(),node.end());
    FOR(i,0,(int)node.size()){
        int id = node[i];
        int v = edge[id].v;
        if(vis[v])  continue;
        res += dfs_21(v,u,1,edge[id].w%m);
        dfs_22(v,u,1,edge[node[i]].w%m);
    }
    return res;
}

//進行分治
void dfs(int u){
    int rt = getroot(u);
    ans += calc(rt);
    vis[rt] = true;
    for(int i = head[rt];i != -1;i = edge[i].nt){
        int v = edge[i].v;
        if(vis[v])  continue;
        dfs(v);
    }
}

void init(){
    memset(head,-1,sizeof(head));
    memset(vis,false,sizeof(vis));
    edge_cnt = 0;
    bit[0] = inv[0] = 1;
    int tem = Inv(10,m);
    FOR(i,1,n+1){
        bit[i] = bit[i-1]*10%m;
        inv[i] = inv[i-1]*tem%m;
    }
    ans = 0;
}


int main(){
    //freopen("test.in","r",stdin);
    bit[0] = inv[0] = 1;
    while(~scanf("%d%d",&n,&m)){
        init();
        FOR(i,1,n){
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            addedge(u,v,w);
            addedge(v,u,w);
        }
        dfs(0);
        printf("%I64d\n",ans);
    }
    return 0;
}
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章