Untitled

user_2740508
plain_text
2 years ago
1.3 kB
5
Indexable
Never
```#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;
}```