CF715C Digit Tree(树上DSU/点分治)

CF715C Digit Tree 树上\(DSU\)做法详解

首先 , \(DSU\)原理应该不必多说 , 即取重儿子 , 将轻儿子合并

统计答案时 , 我们将以 当前所在点\(u\) 为\(LCA\)的点答案进行统计

但是,这题用 \(DSU\) 做的难点在于储存权值

---

我们先暂且不考虑取模

\(DSU\)因为需要从儿子保留，因此容器必须是全局的 , 存的权值也必然是全局的

然而 , 我们所知的全局的值只有两种路径,即\(u\)到\(1\)的路径\(disup\) , 以及\(1\)到\(u\)的路径\(disdown\) , 在维护这两个值时,我们同时维护一个\(dep\)数组(这里的\(dep\)是从\(0\)开始计数的)

我们得到递推式

\[disup[v]=disup[u]+10^{dep[u]}*w \]

\[disdown[v]=10 \cdot disdown[u]+w \]

即将边权为w的边分别接在前后

ps : 因为要多次用到\(10^x\),所以我们先预处理一个数组\(po[x]=10^x\)

所以预处理代码如下

void pre_dfs(int u,int f) {
    for(int i=head[u]; ~i; i=e[i].nxt)if(e[i].to!=f) {
            int v=e[i].to,w=e[i].w;
            disup[v]=disup[u]+po[dep[u]]*w;
            disdown[v]=10*disdown[u]+w;
            dep[v]=dep[u]+1;
            pre_dfs(v,u);
        }
}

维护答案时 , 对于一个点,我们先要将其到\(LCA\)的路径表示出来

设这两段路径为\(up,down\)即

\[up= \frac {disup[u]-disup[lca]}{10^{dep[lca]}} \]

\[down=disdown[u]-disdown[lca]\cdot 10^{dep[u]-dep[lca]} \]

在维护答案时,我们显然是要为一段\(up\)的路径找到一段\(down\)的路径与其匹配

设我们所求的一段路径起始点为\(begin\),结束点为\(end\),则满足条件的路径满足等式

\[up \ \cdot \ 10^{dep[end]-dep[lca]} \ + \ down \equiv \ 0 \pmod{m} \]

我们将\(up\)和\(down\)的表达式带入

\[\frac {(disup[begin]-disup[lca]) \cdot 10^{dep[end]} }{10^{2 \cdot dep[lca]}} + \frac {disdown[end]-disdown[lca]\cdot 10^{dep[end]}} {10^{dep[lca]}} \]

(有那么一点长)

---

现在我们考虑存储权值

对于存入\(up\) , 存入时 , 其实我们只知道 \(disup[begin]\) 这一个值 , 所以我们存下这个值 , 将上式变形

\[disup[begin] = disdown[lca] \cdot 10^{dep[lca]} - \frac{disdown[end] \cdot 10^{2 \cdot dep[lca]} }{10^{dep[end]}}+disup[lca] \]

所以访问答案时 , 我们用上式访问即可

对于存入\(down\),存入时,我们所知的量有\(disdown[end],dep[end]\) , 我们再次变换上式 , 将这两个量提出

\[\frac{disdown[end]}{10^{dep[end]}} = \frac {disdown[lca]}{10^{dep[lca]}} - \frac {disup[begin]-disup[lca]}{10^{2 \cdot dep[lca]}} \]

这两个值都用\(map\)访问即可

现在我们考虑到取模 , 上式中每次运算都加以取模

但是显然模运算之后是不能直接除的 , 需要求出非质数模数的模拟元，需要用扩欧

(这就是为什么题目给出\(gcd(m,10)=1\))

以上是储存的过程

---

还有一个点 , 即如何将轻儿子的值合并

我们要避免儿子自己的信息和自己合并 , 所以要对于每个轻儿子 , 先算一次答案 , 再添加信息

再访问完儿子后 , 再将自己信息加入即可

---

整体代码

#include<cstdio>
#include<vector>
#include<map>
#include<algorithm>
#include<iostream>
#include<cstring>
using namespace std;
const int N=2e5+100;
typedef long long ll;
int n,m;
void Exgcd(ll a,ll b,ll &x,ll &y) {
    if(!b)x=1,y=0;
    else Exgcd(b,a%b,y,x),y-=a/b*x;
}
ll Inv(ll a) {
    ll x,y;
    Exgcd(a,m,x,y);
    x=(x%m+m)%m;
    return x;
}
int po[N]= {1};
struct Edge {
    int to,nxt,w;
} e[N<<1];
int head[N],ecnt;
inline void AddEdge(int u,int v,int w) {
    e[++ecnt].nxt=head[u];
    e[ecnt].to=v;
    e[ecnt].w=w;
    head[u]=ecnt;
}
int cnt[N],son[N],dep[N];
ll disup[N],disdown[N];
void pre_dfs(int u,int f) {
    cnt[u]=1;
    for(int i=head[u]; ~i; i=e[i].nxt)if(e[i].to!=f) {
            int v=e[i].to,w=e[i].w;
            disup[v]=(disup[u]+1ll*po[dep[u]]%m*w)%m;
            disdown[v]=(10ll*disdown[u]%m+w)%m;
            dep[v]=dep[u]+1;
            pre_dfs(v,u);
            cnt[u]+=cnt[v];
            if(cnt[son[u]]<cnt[v])son[u]=v;
        }
}

map<int,int> U,D;
ll ans;
void Calc(int u,int LCA) {
    ll Ub=disup[u],pl=po[dep[LCA]],pe=po[dep[u]],Dl=disdown[LCA],De=disdown[u],Ul=disup[LCA];
    ll x = ((1ll*Dl* pl %m -1ll* De *pl %m* pl %m *Inv(pe) %m+1ll*Ul)%m+m)%m;
    x=(x+m)%m;
    ll y =(1ll* Dl * Inv(pl) %m+1ll*(Ul-Ub)*Inv(pl)%m*Inv(pl)%m)%m;
    y=(y+m)%m;
    ans+=U[x];
    ans+=D[y];
}
void Add(int u) {
    ll x=disup[u],y=1ll*disdown[u]*Inv(po[dep[u]])%m;
    U[x]++,D[y]++;
}
void Upd(int u,int f,int LCA,int k) {
//	cout<<"Now in"<<' '<<u<<' '<<up<<' '<<down<<endl;
//	cout<<"U"<<u<<endl;
    if(k)Calc(u,LCA);
    else Add(u);
//	cout<<"Upd"<<' '<<x<<' '<<y<<' '<<k<<endl;
    for(int i=head[u]; ~i; i=e[i].nxt) {
        int v=e[i].to;
        if(v==f)continue;
        Upd(v,u,LCA,k);
    }
//	Add(u);
//	s+=k*2;
}

void dfs_getans(int u,int f,int h) {
    for(int i=head[u]; ~i; i=e[i].nxt)
        if(e[i].to!=f&&e[i].to!=son[u])
            dfs_getans(e[i].to,u,0);
    if(son[u])dfs_getans(son[u],u,1);
//	int lst=ans;
    for(int i=head[u]; ~i; i=e[i].nxt) if(e[i].to!=f&&e[i].to!=son[u]) {
            Upd(e[i].to,u,u,1);
            Upd(e[i].to,u,u,0);
        }
    Calc(u,u);
    Add(u);
    if(!h)U.clear(),D.clear();
}


int main() {
    scanf("%d%d",&n,&m);
    memset(head,-1,sizeof head);
    for(int i=1; i<=n+3; i++)po[i]=1LL*po[i-1]*10%m;
    for(int i=2,u,v,w; i<=n; i++) {
        scanf("%d%d%d",&u,&v,&w);
        u++,v++;
        AddEdge(u,v,w);
        AddEdge(v,u,w);
    }
//	dep[1]=1;
    pre_dfs(1,0);
    dfs_getans(1,0,1);
    printf("%lld\n",ans);
    return 0;
}