Untitled
unknown
plain_text
a year ago
3.1 kB
6
Indexable
``` #include <bits/stdc++.h> using namespace std; #define rep(i, a, b) for(int i = a; i < (b); ++i) #define all(x) begin(x), end(x) #define sz(x) (int)(x).size() typedef long long ll; typedef pair<int, int> pii; typedef vector<int> vi; const int N = 5001 , M = 2001; int sz[N] , dp[N][M] , val[N]; vector<int> adj[N]; ll sum[N]; const ll mod = (119 << 23) + 1, root = 62; // = 998244353 // For p < 2^30 there is also e.g. 5 << 25, 7 << 26, 479 << 21 // and 483 << 21 (same root). The last two are > 10^9. ll modpow(ll b, ll e) { ll ans = 1; for (; e; b = b * b % mod, e /= 2) if (e & 1) ans = ans * b % mod; return ans; } typedef vector<ll> vl; void ntt(vl &a) { int n = sz(a), L = 31 - __builtin_clz(n); static vl rt(2, 1); for (static int k = 2, s = 2; k < n; k *= 2, s++) { rt.resize(n); ll z[] = {1, modpow(root, mod >> s)}; rep(i,k,2*k) rt[i] = rt[i / 2] * z[i & 1] % mod; } vi rev(n); rep(i,0,n) rev[i] = (rev[i / 2] | (i & 1) << L) / 2; rep(i,0,n) if (i < rev[i]) swap(a[i], a[rev[i]]); for (int k = 1; k < n; k *= 2) for (int i = 0; i < n; i += 2 * k) rep(j,0,k) { ll z = rt[j + k] * a[i + j + k] % mod, &ai = a[i + j]; a[i + j + k] = ai - z + (z > ai ? mod : 0); ai += (ai + z >= mod ? z - mod : z); } } vl conv(const vl &a, const vl &b) { if (a.empty() || b.empty()) return {}; int s = sz(a) + sz(b) - 1, B = 32 - __builtin_clz(s), n = 1 << B; int inv = modpow(n, mod - 2); vl L(a), R(b), out(n); L.resize(n), R.resize(n); ntt(L), ntt(R); rep(i,0,n) out[-i & (n - 1)] = (ll)L[i] * R[i] % mod * inv % mod; ntt(out); return {out.begin(), out.begin() + s}; } void dfs(int u, int p) { sz[u] = 1; dp[u][val[u]] = 1; for (auto &v : adj[u]) if (v != p) { dfs(v, u); vl dp_u(M, 0), dp_v(M, 0); for (int i = 0; i < M; ++i) { dp_u[i] = dp[u][i]; dp_v[i] = dp[v][i]; } vl conv_result = conv(dp_u, dp_v); for (int s = 0; s < M; ++s) { dp[u][s] = (dp[u][s] + conv_result[s]) % mod; } } } int main(){ ios::sync_with_stdio(0); cin.tie(NULL); cout.tie(0); #ifndef ONLINE_JUDGE freopen("input.txt", "r", stdin); freopen("output.txt", "w", stdout); freopen("error.txt", "w", stderr); #endif int x , n , q; cin >> n >> x >> q; --x; for(int i = 0; i < n; ++i) cin >> val[i]; for(int i = 1; i < n; ++i){ int u , v; cin >> u >> v; --u; --v; adj[u].emplace_back(v); adj[v].emplace_back(u); } dfs(x , x); for(int s = 1; s < M; ++s){ sum[s] += sum[s - 1]; if(sum[s] >= mod) sum[s] -= mod; } for(int s = 1; s < M; ++s){ sum[s] = sum[s - 1] + dp[x][s]; if(sum[s] >= mod) sum[s] -= mod; } while(q--){ int l , r; cin >> l >> r; ll ans = (sum[r] - sum[l - 1] + mod) % mod; cout << ans << '\n'; } } ```
Editor is loading...
Leave a Comment