競技プログラミング日記

主に AtCoder の記事です

AtCoder Beginner Contest 222E

ABC222E

\(a\) により path \(p\) は一意に定まる. よって,\(p\) の中で各辺が何回現れるのかをあらかじめ求めておけば, 各辺の色分けを固定したときに赤と青がそれぞれ何回になるか分かる.

赤,青に塗った辺の本数をそれぞれ \(r,b\) とおき, 辺 \(e\) が path \(p\) の中に現れる回数を \(c(e)\) とする. \(s := \sum_{e \in E} c(e)\) とおく.
このとき,\(r + b = s\) かつ \(r - b = k\) となるから, \(r = \frac{s+k}{2}, b = \frac{s-k}{2}\) となる. よって,塗り方が存在するためには \(k+s \equiv 0 \ mod \ 2\) かつ \(k+s \geq 0\) かつ \(k-s \geq 0\) が必要.以下,これらを仮定する.
赤の塗り分け方は,vector \(c\) のいくつかの和であって, \(r\) を作る方法と対応する. よって,これは部分和問題として解ける.

実装
部分和としては \(r\) 以下しか考えなくてよい.

使っている記号,マクロ等 "https://ecsmtlir.hatenablog.com/entry/2022/12/23/131925"

int main() {
  ll n, m, k ;
  cin >> n >> m >> k;
  vll a(m); cinv(a);
  rep(i,m) a[i]--;
  vector<vector<pll>> to(n);
  rep(i,n-1){
    ll x,y; cin >>x >> y;
    --x, --y;
    to[x].push_back({y,i});
    to[y].push_back({x,i});
  }

  vll c(n-1);
  auto dfs = [&](auto dfs, ll cu, ll pa, ll j) -> bool {
    if(cu == a[j+1]) return true;

    for(auto [ne,i]: to[cu]) if(ne!=pa) {
      c[i]++;
      if(dfs(dfs, ne, cu, j)) return true;
      c[i]--;
    }
    return false;
  };

  rep(i,m-1) dfs(dfs, a[i], -1, i);
 
  ll s = 0;
  rep(i,n-1) {
    s += c[i];
  }

  // Remark : k < 0 もあり得る
  ll r = (s+k)/2; // >= 0
  ll b = (s-k)/2; // >= 0
  if((k+s) % 2 != 0 || s-k<0 || s+k<0 ){
    cout << 0 << endl;
    return 0;
  }

  vector<mint> dp(r+1);
  dp[0] = 1;
  rep(i,n-1) drep(x,r+1){
    if(x+c[i] <= r) dp[x+c[i]] += dp[x];
  }

  cout << dp[r].val() << endl;

 
  return 0;
}