Untitled

#include<bits/stdc++.h>
#define lli long long int
#define mod 1000000007

using namespace std;

vector<vector<int> > v(200010), rems(200010);
vector<int> in(200010), out(200010);
int cnt = 0, ans = 0;

int getCnt(vector<int> &v, int a, int b) {
    int ind1 = lower_bound(v.begin(), v.end(), a) - v.begin();
    int ind2 = lower_bound(v.begin(), v.end(), b) - v.begin();
    --ind2;
    return max(0, ind2 - ind1 + 1);
}

void popInOut(int u, int p, int sum, vector<int> &cost, int k) {
    in[u] = cnt;
    sum = (sum + cost[u]) % k;
    rems[sum].push_back(cnt);
    ++cnt;
    for (int & ent : v[u]) {
        if (ent != p) {
            popInOut(ent, u, sum, cost, k);
        }
    }
    out[u] = cnt;
}

void dfs(int u, int p, int req, vector<int> &cost, int k) {
    ans += getCnt(rems[req], in[u], out[u]);
    req = (req + cost[u]) % k;
    for (int &ent : v[u]) {
        if (ent != p) {
            dfs(ent, u, req, cost, k);
        }
    }
}

int countVerticalPaths(vector<int> cost, int n, vector<int> a, vector<int> b, int k) {
    for (int i = 0; i < n; i++) {
        int x = a[i] - 1, y = b[i] - 1;
        v[x].push_back(y);
        v[y].push_back(x);
    }
    popInOut(0, 0, 0, cost, k);
    dfs(0, 0, 0, cost, k);
    return ans;
}