問題概要

深さ $d (\leq 20)$ (ただし頂点の深さを 1 とする) の完全二分木の頂点に，1から$2^{d-1}-1 $ までの整数値を書いていく．次の条件を満たす書き方の通り数を$10^9+7$ で割ったあまりを求めよ．

任意の頂点の組 i, j について，i が j よりも深い場合は i のほうが j よりも書いてある値が大きい．
l が書き込まれた頂点と r が書き込まれた頂点を結ぶパスに含まれる辺の数は k 本である．

解法概要

以下，ソースコードの書き方と合わせるために，頂点の深さを 0 とする．

1つ目の条件から，頂点の深さごとに書き込まれる値は固定される．深さが $d $ であれば $2^d \leq x \leq 2^{d+1}-1 $ となる $x $ が書き込まれる．これを満たさないと，i が j よりも深いのに i のほうが j よりも書いてある値が小さいような頂点の組 $i, j $ を取ることができてしまう．

また，このことから，l が書き込まれる頂点， r が書き込まれる頂点の深さ $d_l, d_r $ もそれぞれ固定できる．

深さ $d_l, d_r $ が分かっているときに頂点間の距離を知りたい場合，他にどんな情報が必要だろうか．たとえ 2 頂点の深さが同じでも，それらの距離は LCA (Lowest Common Ancestor) がどの頂点かによって異なるはずである．2頂点の LCA の深さを$d_\mathit{lc a} $ とすると，l が書き込まれる頂点と r が書き込まれる頂点の間の距離は

$d_l + d_r - 2d_\mathit{lc a} $

となる．この問題の場合は通り数を求めるべき場合の 2 頂点間の距離は k と決まっているので，上式から $d_l, d_r, k $ から $d_\mathit{lc a} $ を求めることができる．すると，距離が k，$l,r $ が書き込まれる頂点の深さが $d_l, d_r $ となるような $l,r $ の書き込み方の通り数がわかる．

たとえば，$d_l \neq d_\mathit{lc a} $ の場合，

ある1つの頂点を LCA として注目し，$l$ の書き込まれる頂点が LCA の左側の子またはその子孫にあると仮定すると，
- LCA の左側の子またはその子孫で，かつ LCA から見たとき深さが $d_l - d_\mathit{lc a} $ の頂点の数は $2^{d_l - d_\mathit{lc a}-1} $
- r の書き込まれる頂点が LCA の右側の子またはその子孫にあるはずであるから，LCA の右側の子またはその子孫で，かつ LCA から見たとき深さが $d_r - d_\mathit{lc a} $ の頂点の数を求めると， $2^{d_r - d_\mathit{lc a}-1} $
左右反転しても良いので 2 を掛ける
同じ深さの違うノードに LCA を変えた場合を考慮して $2^{d_\mathit{lc a}}$

のように $l,r $ の書き込み方の通り数が求まる．

$d_l = d_\mathit{lc a} $ の場合は，

LCA から見た深さが $d_r $ の頂点の数は $2^{d_r - d_\mathit{lc a}} $
同じ深さの違うノードに LCA を変えた場合を考慮して $2^{d_\mathit{lc a}}$

だけ計算すれば良い．

残りの数字は深さごとにまだ値が書き込まれていない頂点の数 $n $ が決まっているから，深さごとに $n! $ をかけていけばよい．

ソースコード

#include <iostream>
#include <algorithm>

using namespace std;

typedef long long ll;

constexpr ll MOD = 1000000007;
constexpr int N_MAX = 1 << 20;

ll fact[N_MAX+1];

void init(ll n){
    fact[0] = fact[1] = 1;
    for(int i=2;i<=n;++i) {
        fact[i] = (fact[i-1] * (ll)i) % MOD;
        ll k = MOD-2;
        ll a = fact[i];a
        while(k > 0){
            a *= a;
            a %= MOD;
            k  >>= 1;
        }
    }
}

int main() {
    ll d, l, r, k;
    cin >> d >> l >> r >> k;

    init(1LL << d);

    ll ld = 0, rd = 0;
    while((1LL << (ld+1)) <= l) ++ld;
    while((1LL << (rd+1)) <= r) ++rd;

    if(ld+rd > 2*d) {
        cout << 0 << endl;
        return 0;
    }

    ll lca_d = 0;
    while(ld+rd - 2*lca_d > k && lca_d <= min(ld, rd)) {
        ++lca_d;
    }

    if(ld+rd - 2*lca_d != k) {
        cout << 0 << endl;
        return 0;
    }

    ll ans = 0;
    // l と r の置き方
    if(ld == lca_d) {
        ans = (((1LL << (rd-ld)) % MOD) * ((1LL << ld)) % MOD) % MOD;
    } else {
        ans = (((1LL << (rd-lca_d-1)) % MOD) * ((1LL << (ld-lca_d-1))) % MOD) % MOD;
        ans = (ans * ((1LL << lca_d)) % MOD) % MOD;
        (ans *= 2LL) %= MOD;
    }

    // その他の数字の置き方
    for(int i=0;i<d;++i) {
        ll n = 1LL << i;
        if(ld == i) --n;
        if(rd == i) --n;
        ans *= fact[n];
        ans %= MOD;
    }

    cout << ans << endl;
     
}

所感

最近解いた数式立てていく系の問題で一番面白かったかも．

NoiminのNoise

yukicoder No.916 Encounter On A Tree

問題概要

解法概要

ソースコード

所感