Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include<bits/stdc++.h>
- using namespace std;
- #define int long long
- #define double long double
- const int MOD=998244353;
- const int INF=1e18;
- int bin_pow(int a, int b){
- int ans=1;
- while(b){
- if(b&1){
- ans*=a;
- ans%=MOD;
- }
- a*=a;
- a%=MOD;
- b>>=1;
- }
- return ans;
- }
- void dfs(int x, int par, vector<int> g[], vector<vector<int>>&pf, vector<int>&f, int m){
- pf[x][f[x]-1]++;
- for(auto it:g[x]){
- if(it==par) continue;
- dfs(it,x,g,pf,f,m);
- for(int i=0; i<m; i++) pf[x][i]+=pf[it][i];
- }
- }
- void dfs2(int x, int par, vector<int> g[], vector<vector<vector<int>>>&dp, int m, vector<vector<int>>&pf){
- for(int i=0; i<m; i++){
- for(int j=0; j<3; j++){
- dp[x][i][j]+=dp[par][i][j];
- dp[x][i][j]%=MOD;
- }
- if(pf[x][i]==0){
- dp[x][i][2]+=1;
- dp[x][i][2]%=MOD;
- }
- else if(pf[x][i]==1){
- dp[x][i][1]+=1;
- dp[x][i][1]%=MOD;
- dp[x][i][2]+=bin_pow(2,MOD-2);
- dp[x][i][2]%=MOD;
- }
- else{
- dp[x][i][0]+=bin_pow(pf[x][i]-1,MOD-2);
- dp[x][i][0]%=MOD;
- dp[x][i][1]+=bin_pow(pf[x][i],MOD-2);
- dp[x][i][1]%=MOD;
- dp[x][i][2]+=bin_pow(pf[x][i]+1,MOD-2);
- dp[x][i][2]%=MOD;
- }
- }
- for(auto it:g[x]){
- if(it==par) continue;
- dfs2(it,x,g,dp,m,pf);
- }
- }
- signed main() {
- ios_base::sync_with_stdio(false);
- cin.tie(NULL);
- int t;
- cin>>t;
- while(t--){
- int n, m;
- cin>>n>>m;
- vector<int> v(m);
- vector<int> f(n+1);
- for(int i=1; i<=n; i++) cin>>f[i];
- for(int i=0; i<m; i++) cin>>v[i];
- vector<int> g[n+1];
- for(int i=0; i<n-1; i++){
- int a, b;
- cin>>a>>b;
- g[a].push_back(b);
- g[b].push_back(a);
- }
- vector<vector<int>> pf(n+1, vector<int>(m,0));
- dfs(1,0,g,pf,f,m);
- double ini = 0;
- int init = 0;
- for(int i=1; i<=n; i++){
- for(int j=0; j<m; j++){
- if(!pf[i][j]) continue;
- int dl=v[j]*bin_pow(pf[i][j],MOD-2);
- dl%=MOD;
- init+=dl;
- init%=MOD;
- }
- }
- vector<vector<vector<int>>> dp(n+1, vector<vector<int>>(m, vector<int>(3,0)));
- dfs2(1,0,g,dp,m,pf);
- int q;
- cin>>q;
- int ans = init;
- while(q--){
- int x, r;
- cin>>x>>r;
- if(f[x]==r) cout<<ans<<'\n';
- else{
- int d=f[x]-1;
- r--;
- int tm=ans;
- tm-=dp[x][d][1]*v[d];
- tm%=MOD;
- tm-=dp[x][r][1]*v[r];
- tm%=MOD;
- tm+=dp[x][d][0]*v[d];
- tm%=MOD;
- tm+=dp[x][r][2]*v[r];
- tm%=MOD;
- tm+=MOD;
- tm%=MOD;
- cout<<tm<<'\n';
- }
- }
- }
- }
Advertisement
Add Comment
Please, Sign In to add comment