競技プログラミング日記

主に AtCoder の記事です

AtCoder Beginner Contest 321E問題

\(\def \set #1#2{\{ #1 \ \vert \ #2 \}}\)

ABC321E

解法

1-indexed で complete binary tree の頂点に番号を付ける. LCA \(v\) を全探索. \(v\) を root とする subtree において,\(v\) からの距離が \(d\) である頂点全体の個数 \(f(v,d)\) を求める.
頂点 \(i \neq 1\) に対して,\(i\) と同じ parent を持つ頂点は \(i xor 1\) と書ける.
頂点 \(x\) に対して,root の方に向かって移動しながら \(f(x,*)\) を求める.

注意

\(x\) が root 1 に到達する前に終わることがある.

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

int main() {
  ll t; cin >> t;
 
  rep(ti,t){
    ll n,x,k;
    cin >> n >> x >> k;

    auto f = [&](ll v, ll d) -> ll {
      ll l = v, r = v;
      if(l > n) return 0;

      rep(i,d){
        l <<= 1;
        r = (r << 1) | 1;
        if(l > n) return 0;
      }

      chmin(r,n);
      return r-l+1;
    };

    ll ans = 0;
    ans += f(x,k);
    while(x > 1 && k >= 2){
      // [*] -------------------------------
        // ans += f(x^1, k-2);
        ans += f(x>>1, k-1);
        ans -= f(x, k-2);
      //---------------------------------
      x >>= 1;
      k--;
    }

    // k == 0: 数える必要なし
    // k == 1:
    //    x > 1 だけ数える.
    //    x == 1 のときは [*]で既に数えている.
    if(x > 1 && k == 1){
      ans++;
    }

    cout << ans << endl;

  }

  return 0;
}