Submission #52926899

---

Source Code Expand

#include <bits/stdc++.h>
#include <atcoder/modint>
const int MOD = 998244353;
using mint = atcoder::static_modint<MOD>;
using namespace std;

typedef long long ll;
#define rep(i, a, n) for (ll i = a; i < (ll)(n); i++)
#define per(i, a, n) for (ll i = n; i-- > (ll)(a);)

ll read() { ll r; scanf("%lld", &r); return r; }

#define N 200010
int a[N];
vector<int> e[N];

int sz[N],fa[N],dep[N];
void dfs1(int u,int F) {
  per(i,0,size(e[u])) if(e[u][i]==F) {
    swap(e[u][i],e[u].back());
    e[u].pop_back();
  }
  sz[u]=1;
  dep[u]=dep[fa[u]=F]+1;
  rep(i,0,size(e[u])) {
    int v=e[u][i];
    dfs1(v,u);
    sz[u]+=sz[v];
    if(sz[v]>sz[e[u][0]]) swap(e[u][0],e[u][i]); // 重儿子, size
  }
}

int dfn[N],root[N];
void dfs2(int u,int R,int &cur) {
  dfn[u]=cur++; // dfs number
  root[u]=R;
  if(sz[u] == 1) return; // 叶子
  dfs2(e[u][0],R,cur); // 重儿子
  rep(i,1,size(e[u])) dfs2(e[u][i],e[u][i],cur); // 轻儿子
}
int lca(int u,int v) {
  for(;root[u]!=root[v];u=fa[root[u]]) if(dep[root[u]]<dep[root[v]]) swap(u,v);
  return dep[u]<dep[v]?u:v; // u,v在一条重链上
}
vector<int> eg[N];
vector<int> fake(vector<int> v) {
  auto cmp=[&](int x,int y){return dfn[x]<dfn[y];};
  vector<int> ans=v;
  sort(v.begin(),v.end(),cmp); // 按照 dfs序 sort
  rep(i,1,size(v)) ans.push_back(lca(v[i-1],v[i])); // dfn相邻的lca
  ans.push_back(1); // 根节点, 加了简化处理不影响结果
  sort(ans.begin(),ans.end(),cmp);
  ans.resize(unique(ans.begin(),ans.end())-ans.begin()); // 去重
  rep(i,1,size(ans)) eg[lca(ans[i-1],ans[i])].push_back(ans[i]);
  return ans;
}
mint f[N];
void dfs(int x,int C,mint&ans) {
  f[x]=1;
  mint tot=0; // 单儿子的方案和, 在根非颜色时 不贡献
  for(auto i:eg[x]) {
    dfs(i,C,ans);
    tot+=f[i];
    f[x]*=(f[i]+1);
  }
  if(a[x]!=C){
    f[x]-=1;
    ans+=f[x]-tot;
  }else{
    ans+=f[x];
  }
}
vector<int> c2i[N]; // color to index
int main() {
  int n=read();
  rep(i,1,n+1) {
    a[i]=read();
    c2i[a[i]].push_back(i);
  }
  rep(i,1,n) {
    int u=read();
    int v=read();
    e[u].push_back(v);
    e[v].push_back(u);
  }
  dfs1(1,0); // 父节点, 子树大小, 深度, 重儿子
  int cur=0;
  dfs2(1,1,cur); // 处理重链
  mint ans=0;
  rep(i,1,n+1) if(c2i[i].size()) {
    vector<int> v=fake(c2i[i]);
    dfs(1,i,ans);
    for(auto j:v) eg[j].clear();
  }
  printf("%d",ans.val());
  return 0;
}

Submission Info

Compile Error

Main.cpp: In function ‘ll read()’:
Main.cpp:11:24: warning: ignoring return value of ‘int scanf(const char*, ...)’ declared with attribute ‘warn_unused_result’ [-Wunused-result]
   11 | ll read() { ll r; scanf("%lld", &r); return r; }
      |                   ~~~~~^~~~~~~~~~~~

Judge Result