In this HackerEarth Little Shino and Path Divisor problem solution Given a weighted undirected tree of N nodes and Q queries. Each query contains an integer D. For each query, find the number of paths in the given tree such that all the edges in the path are divisible by D.


HackerEarth Little Shino and Path Divisor problem solution


HackerEarth Little Shino and Path Divisor problem solution.

#include <bits/stdc++.h>

#define ll long long
#define ull unsigned long long
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define be begin()
#define en end()
#define all(x) (x).begin(),(x).end()
#define alli(a, n, k) (a+k),(a+n+k)
#define REP(i, a, b, k) for(__typeof(a) i = a;i < b;i += k)
#define REPI(i, a, b, k) for(__typeof(a) i = a;i > b;i -= k)
#define REPITER(it, a) for(__typeof(a.begin()) it = a.begin();it != a.end(); ++it)

#define y0 sdkfaslhagaklsldk
#define y1 aasdfasdfasdf
#define yn askfhwqriuperikldjk
#define j1 assdgsdgasghsf
#define tm sdfjahlfasfh
#define lr asgasgash
#define norm asdfasdgasdgsd
#define have adsgagshdshfhds

#define eps 1e-6
#define pi 3.141592653589793

using namespace std;

template<class T> inline T gcd(T x, T y) { if (!y) return x; return gcd(y, x%y); }
template<class T> inline T mod(T x) { if(x < 0) return -x; else return x; }

typedef vector<int> VII;
typedef vector<ll> VLL;
typedef pair<int, int> PII;
typedef pair<int, PII > PPII;
typedef vector< PII > VPII;
typedef vector< PPII > VPPI;

const int MOD = 1e9 + 7;
const int INF = 1e9;
const int MAX = 1e5 + 5;
VPII edges[MAX];

int id[MAX], sz[MAX];
ll ANS[MAX], ans;

void init(int n)
{
    REP(i, 0, n+1, 1)
        id[i] = i, sz[i] = 1;
}

int root(int x)
{
    while(x != id[x])
    {
        id[x] = id[id[x]];
        x = id[x];
    }
    return x;
}

void union1(int x, int y)
{
    int p = root(x);
    int q = root(y);
    if(p != q)
    {
        ans += ((ll)sz[p]*(ll)sz[q]);
        if(sz[p] < sz[q])
            id[p] = q, sz[q] += sz[p];
        else
            id[q] = p, sz[p] += sz[q];
    }
}

int main(int argc, char* argv[])
{
    if(argc == 2 or argc == 3) freopen(argv[1], "r", stdin);
    if(argc == 3) freopen(argv[2], "w", stdout);
    ios::sync_with_stdio(false);
    int n, q, d, x, y, w, k;
    PII p;

    // Input

    cin >> n >> q;

    REP(i, 1, n, 1)
    {
        cin >> x >> y >> w;
        k = 1;
        while(k*k <= w)
        {
            if(w % k == 0)
            {
                if(k*k == w)
                    edges[k].pb(mp(x, y));
                else
                {
                    edges[k].pb(mp(x, y));
                    edges[w / k].pb(mp(x, y));
                }
            }
            k++;
        }
    }

    // Solution

    init(MAX);

    REP(i, 1, MAX, 1)
    {
        ans = 0;
        REP(j, 0, edges[i].size(), 1)
        {
            p = edges[i][j];
            union1(p.fi, p.se);
        }
        ANS[i] = ans;
        REP(j, 0, edges[i].size(), 1)
        {
            p = edges[i][j];
            id[p.fi] = p.fi;
            sz[p.fi] = 1;
            id[p.se] = p.se;
            sz[p.se] = 1;
        }
    }

    // Query / Output
    
    REP(cases, 1, q+1, 1)
    {
        cin >> d;
        cout << ANS[d] << endl;
    }
    return 0;
}