?
# | Author | Problem | Lang | Verdict | Time | Memory | Sent | Judged | |
---|---|---|---|---|---|---|---|---|---|
238715837 |
Practice: Schucking_Sattin |
434E - 23 | C++14 (GCC 6-32) | Accepted | 1247 ms | 18372 KB | 2023-12-24 18:38:52 | 2023-12-24 18:38:54 |
#include<bits/stdc++.h> #define LL long long #define DB double int MOD; #define ls(x) x << 1 #define rs(x) x << 1 | 1 #define lowbit(x) x & (-x) #define PII pair<int, int> #define MP make_pair #define VI vector<int> #define VII vector<int>::iterator #define all(x) x.begin(), x.end() #define EB emplace_back #define SI set<int> #define SII set<int>::iterator #define QI queue<int> #define fi first #define se second using namespace std; template<typename T> void chkmn(T &a, const T &b) { (a > b) && (a = b); } template<typename T> void chkmx(T &a, const T &b) { (a < b) && (a = b); } int inc(const int &a, const int &b) { return a + b >= MOD ? a + b - MOD : a + b; } int dec(const int &a, const int &b) { return a - b < 0 ? a - b + MOD : a - b; } int mul(const int &a, const int &b) { return 1LL * a * b % MOD; } int sqr(const int &a) { return 1LL * a * a % MOD; } void Inc(int &a, const int &b) { ((a += b) >= MOD) && (a -= MOD); } void Dec(int &a, const int &b) { ((a -= b) < 0) && (a += MOD); } void Mul(int &a, const int &b) { a = 1LL * a * b % MOD; } void Sqr(int &a) { a = 1LL * a * a % MOD; } int qwqmi(int x, int k = MOD - 2) { int res = 1; while(k) { if(k & 1) Mul(res, x); Sqr(x), k >>= 1; } return res; } template<typename T> void read(T &x) { x = 0; int f = 1; char ch = getchar(); while(!isdigit(ch)) { if(ch == '-') f = -1; ch = getchar(); } while(isdigit(ch)) { x = (x << 1) + (x << 3) + (ch ^ 48); ch = getchar(); } x= x * f; } const int N = 1e5 + 5; const int INF = 1e9 + 7; int n, m, K, invK, w[N]; int qwq[N], qwqinv[N]; int a[N], b[N]; vector<int> G[N]; struct Shiroha { int root; int nowsz; int maxsz; int sz[N]; int vis[N]; unordered_map<int, int> mp1, mp2; void get_root(int u, int fa) { sz[u] = 1; int maxx = 0; for(auto v : G[u]) { if(v == fa || vis[v]) continue; get_root(v, u); sz[u] += sz[v]; chkmx(maxx, sz[v]); } chkmx(maxx, nowsz - sz[u]); if(maxx < maxsz) { maxsz = maxx; root = u; } } void modify(int u, int fa, int w1, int w2, int dep, int op) { mp1[w1] += op; mp2[mul(dec(m, w2), qwqinv[dep + 1])] += op; for(auto v : G[u]) { if(v == fa || vis[v]) continue; modify(v, u, inc(w1, mul(w[v], qwq[dep + 1])), inc(mul(w2, K), w[v]), dep + 1, op); } } void dfs(int u, int fa, int w1, int w2, int dep) { a[u] += mp2[w1]; b[u] += mp1[mul(dec(m, w2), qwqinv[dep + 1])]; for(auto v : G[u]) { if(v == fa || vis[v]) continue; dfs(v, u, inc(w1, mul(w[v], qwq[dep + 1])), inc(mul(w2, K), w[v]), dep + 1); } } void calc(int u) { // mp1.clear(); // mp2.clear(); while(!mp1.empty()) mp1.erase(mp1.begin()); while(!mp2.empty()) mp2.erase(mp2.begin()); modify(u, 0, w[u], w[u], 0, 1); a[u] += mp2[0]; b[u] += mp1[m]; for(auto v : G[u]) { if(vis[v]) continue; modify(v, u, inc(w[u], mul(w[v], K)), inc(mul(w[u], K), w[v]), 1, -1); dfs(v, u, w[v], w[v], 0); modify(v, u, inc(w[u], mul(w[v], K)), inc(mul(w[u], K), w[v]), 1, 1); } } void PDC(int u, int dep) { // cerr << dep << '\n'; vis[u] = 1, calc(u); for(auto v : G[u]) { if(vis[v]) continue; get_root(v, u); nowsz = sz[v]; maxsz = INF; get_root(v, u); PDC(root, dep + 1); } } }umi; LL ans; int main() { read(n), read(MOD), read(K), read(m); for(int i = 1; i <= n; ++i) read(w[i]); for(int i = 1; i < n; ++i) { int x, y; read(x), read(y); G[x].EB(y), G[y].EB(x); } qwq[0] = qwqinv[0] = 1; invK = qwqmi(K); for(int i = 1; i < N; ++i) { qwq[i] = mul(qwq[i - 1], K); qwqinv[i] = mul(qwqinv[i - 1], invK); } umi.PDC(1, 0); ans = 1LL * n * n * n; for(int i = 1; i <= n; ++i) { ans -= 3LL * n * a[i]; ans += 1LL * a[i] * b[i]; ans += 1LL * a[i] * a[i]; ans += 1LL * b[i] * b[i]; } cout << ans << '\n'; return 0; }
?
?
?
?