2019(dp/点分治)

2019

题意：

输入$n(2e4)$
接下来$n-1$行，每行$u,v,w(<2019)$表示树边
求点对个数，满足两点路径距离是2019的倍数。
多组数据$\sum n\le 1e5$

题解1(dp)：

$dp[u][j]$表示其中一点为$u$,另一点为子树一点$v$,距离为$j$出现的个数。
在转移的同时，经过这点的路径的贡献算上。
时间复杂度$O(nk)$

代码：

#include<bits/stdc++.h>
using namespace std;
const int N=2e4+9;
int n;
int dp[N][2029];
struct Edge{
    int v,w,nxt;
}e[N<<1];
int head[N],cnt;
inline void add(int u,int v,int w){
    e[cnt]=(Edge){v,w,head[u]};
    head[u]=cnt++;
}
int ans;
void solve(int u,int fa){
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(e1.v==fa)continue;
        solve(e1.v,u);
        for(int j=0;j<2019;j++){
            int t=(j+e1.w)%2019;
            ans+=((t==0)+dp[u][(2019-t)%2019])*dp[e1.v][j];
        }
        for(int j=0;j<2019;j++){
            int t=(j+e1.w)%2019;
            dp[u][t]+=dp[e1.v][j];
        }
    }
    dp[u][0]++;
}
int main(){
  //  freopen("tt.in","r",stdin),freopen("tt.out","w",stdout);
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    while(cin>>n){
        for(int i=1;i<=n;i++)head[i]=-1;cnt=0;ans=0;
        for(int i=1;i<=n;i++)for(int j=0;j<2019;j++)dp[i][j]=0;
        for(int i=1,u,v,w;i<n;i++){
            cin>>u>>v>>w;
            add(u,v,w),add(v,u,w);
        }
        solve(1,0);
        cout<<ans<<endl;
    }
    return 0;
}

题解2(点分治)

就是个板子，没啥好说的。
时间复杂度$O(2019n)$

代码：

#include<bits/stdc++.h>
using namespace std;
const int N=2e4+9;
int n,ans;
struct Edge{
    int v,w,nxt;
}e[N<<1];
int head[N],cnt;
inline void add(int u,int v,int w){
    e[cnt]=(Edge){v,w,head[u]};
    head[u]=cnt++;
}
int rt,vis[N];
int siz[N],mx[N]={N},tot;
void getroot(int u,int fa){
    siz[u]=1;mx[u]=0;
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(vis[e1.v]||e1.v==fa)continue;
        getroot(e1.v,u);
        siz[u]+=siz[e1.v];
        mx[u]=max(mx[u],siz[e1.v]);
    }
    mx[u]=max(mx[u],tot-siz[u]);
    if(mx[rt]>mx[u])rt=u;
}
int he[2029],tmp[2029];
void dfs(int u,int fa,int val){
    tmp[val]++;
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(e1.v==fa||vis[e1.v])continue;
        dfs(e1.v,u,(val+e1.w)%2019);
    }
}
void calc(int u,int val){
    memset(tmp,0,sizeof(tmp));
    dfs(u,0,val);
    ans+=tmp[0]*(he[0]+1);
    for(int i=1;i<2019;i++)ans+=tmp[i]*he[2019-i];
    for(int i=0;i<2019;i++)he[i]+=tmp[i];
}
void solve(int u){
  //  cout<<"rt "<<u<<endl;
    memset(he,0,sizeof(he));
    vis[u]=1;
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(vis[e1.v])continue;
        calc(e1.v,e1.w);
    }
    for(int i=head[u];~i;i=e[i].nxt){
        Edge e1=e[i];
        if(vis[e1.v])continue;
        rt=0;tot=siz[e1.v];getroot(e1.v,u),solve(rt);
    }
}
int main(){
 //   freopen("tt.in","r",stdin),freopen("tt.out","w",stdout);
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    while(cin>>n){
        for(int i=1;i<=n;i++)head[i]=-1,vis[i]=0;cnt=0;ans=0;
        for(int i=1,u,v,w;i<n;i++){
            cin>>u>>v>>w;
            add(u,v,w),add(v,u,w);
        }
        rt=0,tot=n;getroot(1,0);
        solve(rt);
        cout<<ans<<endl;
    }
    return 0;
}