Codeforces 716E E. Digit Tree

E. Digit Tree
time limit per test3 seconds
memory limit per test256 megabytes
inputstandard input
outputstandard output
ZS the Coder has a large tree. It can be represented as an undirected connected graph of n vertices numbered from 0 to n - 1 and n - 1 edges between them. There is a single nonzero digit written on each edge.

One day, ZS the Coder was bored and decided to investigate some properties of the tree. He chose a positive integer M, which is coprime to 10, i.e. .

ZS consider an ordered pair of distinct vertices (u, v) interesting when if he would follow the shortest path from vertex u to vertex v and write down all the digits he encounters on his path in the same order, he will get a decimal representaion of an integer divisible by M.

Formally, ZS consider an ordered pair of distinct vertices (u, v) interesting if the following states true:

Let a1 = u, a2, …, ak = v be the sequence of vertices on the shortest path from u to v in the order of encountering them;
Let di (1 ≤ i < k) be the digit written on the edge between vertices ai and ai + 1;
The integer is divisible by M.
Help ZS the Coder find the number of interesting pairs!

Input
The first line of the input contains two integers, n and M (2 ≤ n ≤ 100 000, 1 ≤ M ≤ 109, ) — the number of vertices and the number ZS has chosen respectively.

The next n - 1 lines contain three integers each. i-th of them contains ui, vi and wi, denoting an edge between vertices ui and vi with digit wi written on it (0 ≤ ui, vi < n,  1 ≤ wi ≤ 9).

Output
Print a single integer — the number of interesting (by ZS the Coder’s consideration) pairs.

Examples
input
6 7
0 1 2
4 2 4
2 0 1
3 0 9
2 5 7
output
7
input
5 11
1 2 3
2 0 3
3 0 3
4 3 3
output
8
Note
In the first sample case, the interesting pairs are (0, 4), (1, 2), (1, 5), (3, 2), (2, 5), (5, 2), (3, 5). The numbers that are formed by these pairs are 14, 21, 217, 91, 7, 7, 917 respectively, which are all multiples of 7. Note that (2, 5) and (5, 2) are considered different.

In the second sample case, the interesting pairs are (4, 0), (0, 4), (3, 2), (2, 3), (0, 1), (1, 0), (4, 1), (1, 4), and 6 of these pairs give the number 33 while 2 of them give the number 3333, which are all multiples of 11.

#include<cstdio>
#include<algorithm>
#include<iostream>
#include<cstring>
#include<ctime>
#include<map>
using namespace std;
#define ll long long
#define L 1048576
#define MAXN 100005
map<int,int> A;
char _buf[L],*S,*T,c;
char gc()
{
    if(S==T)
    {
        T=(S=_buf)+fread(_buf,1,L,stdin);
        if(S==T)return 0;
    }
    return *S++;
}
void get(int& x)
{
    for(c=gc();c<'0'||c>'9';c=gc());
    for(x=c^'0',c=gc();c>='0'&&c<='9';c=gc())x=(x<<3)+(x<<1)+(c^'0');
}
ll ans;
int n,N,M,I[MAXN],t[MAXN],m,i,j,k,l,h[MAXN],ne[MAXN<<1],p[MAXN<<1],w[MAXN<<1],R,s[MAXN],f[MAXN],a[MAXN],b[MAXN],sum;
bool v[MAXN];
void dfs(int now,int fa)
{
    s[now]=1;
    f[now]=0;
    for(int i=h[now];i;i=ne[i])if(!v[p[i]]&&p[i]!=fa)
    {
        dfs(p[i],now);
        s[now]+=s[p[i]];
        f[now]=max(f[now],s[p[i]]);
    }
    f[now]=max(f[now],sum-s[now]);
    if(f[now]<f[R])R=now;
}
void dfs(int now,int d,int d1,int D,int fa)
{
    a[++M]=d;
    b[M]=d1;
    for(int i=h[now];i;i=ne[i])if(!v[p[i]]&&p[i]!=fa)dfs(p[i],(d+(ll)w[i]*t[D])%N,(d1+(ll)w[i]*I[D+1])%N,D+1,now);
}
void work(int now,int d,int d1,int D,int x)
{
    M=0;
    dfs(now,d,d1,D,0);
    A.clear();
    int i;
    for(i=1;i<=M;i++)A[a[i]]++;
    for(i=1;i<=M;i++)ans+=x*A[(N-b[i])%N];
}
void work(int now)
{
   v[now]=1;
    work(now,0,0,0,1);
   for(int i=h[now];i;i=ne[i])if(!v[p[i]])
   {
       work(p[i],w[i]%N,(ll)w[i]*I[1]%N,1,-1);
      R=0;
      sum=s[p[i]];
      dfs(p[i],0);
      work(R);
   }
}
int main()
{
    get(n);
    get(N);
    for(i=1;i<10;i++)if(((ll)i*N+1)%10==0)break;
    for(I[0]=1,I[1]=((ll)i*N+1)/10,i=2;i<=n;i++)I[i]=(ll)I[i-1]*I[1]%N;
    for(t[0]=i=1;i<=n;i++)t[i]=(ll)t[i-1]*10%N;
   for(i=1;i<n;i++)
   {
       get(j);
       get(k);
       get(l);
       j++;
       k++;
       p[++m]=k;
       w[m]=l;
      ne[m]=h[j];
      h[j]=m;
      p[++m]=j;
      w[m]=l;
      ne[m]=h[k];
      h[k]=m;
    }
    f[0]=n+1;
    sum=n;
    R=0;
    dfs(1,0);
    work(R);
    cout<<ans-n<<endl;
    return 0;
}