AGC060 C: Large Heap - SSRS の競プロ日記

atcoder.jp

解法

[1] 条件の言い換え

対称性により、不等号の向きをすべて反転させても求める確率は変わりません。後の計算がやりやすくなるので、不等号の向きをすべて反転させて考えます。

長さ $2^N-1$ の実数列 $X = (x_1, x_2, \dots, x_{2^N-1})$ が以下の条件を満たすとき、 $X$ は $[0,1]$ 上の列であると呼ぶことにします。

$0 \leq x_i \leq 1 \ (1 \leq i \leq 2^N-1)$

$[0,1]$ 上の列 $X = (x_1, x_2, \dots, x_{2^N-1})$ がさらに以下の条件をすべて満たすとき、 $X$ は $[0,1]$ 上のヒープ的な列であると呼ぶことにします。

$x_i > x_{2i} \ (1 \leq i \leq 2^{N-1}-1)$
$x_i > x_{2i+1} \ (1 \leq i \leq 2^{N-1}-1)$

このとき、求める確率は $[0,1]$ 上のヒープ的な列 $X$ を一様ランダムに $1$ つ選んだときに $x_U > x_V$ となる確率に一致します。これは以下のようにしてわかります。

$[0,1]$ 上の列 $(x_1, x_2, \dots, x_{2^N-1})$ がヒープ的であるか、また $x_U > x_V$ であるかは、 $x_i$ の間の大小関係のみに依存する。
$[0,1]$ 上の列を一様ランダムに $1$ つ選んだとき、等しい要素が存在する確率は $0$ であるから無視できる。また、各要素が相異なるとき、要素間の大小関係としてありうるものは $(2^N-1)!$ 通りあるが、それらは等確率で現れる。

さらに、以下のような根付き木 $T$ を考えます。

$2^N-1$ 個の頂点を持つ。
根は頂点 $1$ である。
$1 \leq i \leq 2^{N-1}-1$ について、頂点 $i$ は頂点 $2i, 2i+1$ の $2$ つの子を持つ。

このとき、 $[0,1]$ 上の列 $(x_1, x_2, \dots, x_{2^N-1})$ がヒープ的であることは、木 $T$ の各頂点 $i$ に $x_i$ を書いたとき、頂点 $u$ が頂点 $v$ の親であるような木 $T$ の頂点の組 $(u,v)$ すべてについて頂点 $u$ に書かれた値が頂点 $v$ に書かれた値より大きいことと同値になります。

以上により、順列に関する問題を、木の各頂点に区間 $[0,1]$ に含まれる実数を書く問題に言い換えることができました。

$[0,1]$ 上の列 $X = (x_1, x_2, \dots, x_{2^N-1})$ を一様ランダムに選んだとき、 $X$ がヒープ的である確率を $p_1(N)$ , $X$ がヒープ的でかつ $x_U > x_V$ である確率を $p_2(N, A, B)$ とおくと、答えは $\displaystyle\frac{p_2(N, A, B)}{p_1(N)}$ となります。よって、 $p_1(N), p_2(N, A, B)$ をそれぞれ求めればよいです。

[2] $p_1(N)$ の計算

$p_1(N)$ を求めるときは、実数で考えるのではなく、順列のままで考えると求めやすいです。

以下の事実が典型として知られています。

$N$ 頂点の根付き木がある。長さ $N$ の順列 $P_1, P_2, \dots, P_N$ を一様ランダムに $1$ つ選んだとき、根以外の全ての頂点 $v$ に対し、その親を $p$ とすると $P_p > P_v$ となる確率は、頂点 $i$ の部分木の頂点数を $\text{size}_i$ とおくと $\displaystyle \frac{1}{\prod_{i=1}^N \text{size}_i}$ である。

木 $T$ は深さ $N-1$ の完全二分木です。 $i = 0, 1, \dots, N-1$ に対し、深さ $i$ の頂点は $2^i$ 個存在し、それらの頂点の部分木の頂点数はすべて $2^{N-i}-1$ です。したがって、 $\displaystyle p_1(N) = \frac{1}{\prod_{i=0}^{N-1} \left(2^{N-i}-1\right)^{2^i}}$ となります。

もしくは、最終的に用いる値が $\displaystyle\frac{1}{p_1(N)}$ であることから、漸化式 $\displaystyle\frac{1}{p_1(1)}=1, \frac{1}{p_1(N)}=\left(\frac{1}{p_1(N-1)}\right)^2\left(2^N-1\right)$ を用いると簡単に計算することができます。

[3] $p_2(N, A, B)$ の計算

頂点 $U, V$ に書く値については特別な制約があるため、最初に固定しておきます。具体的には、頂点 $U$ には必ず $a$ を、頂点 $V$ には必ず $b$ を書くことにします。ただし、 $0 \leq b \lt a \leq 1$ とします。

以下のような木 DP を考えます。

$\text{dp}[v](a,b,x)$ : 木 $T$ の頂点 $v$ の部分木の各頂点に区間 $[0,1]$ に含まれる実数を独立に一様ランダムに書くが、頂点 $U$ に数を書く場合には必ず $a$ を書き、頂点 $V$ に数を書く場合には必ず $b$ を書くとき、頂点 $v$ に書いた数が $x$ 以下であり、かつ頂点 $v$ の部分木で大小関係の条件が満たされている確率

ただし、頂点 $v$ の部分木で大小関係の条件が満たされているとは、頂点 $v$ の部分木に含まれる各頂点 $w$ に対し頂点 $w$ に書いた数が頂点 $w$ の子に書いた数より大きくなることと定義します。また、 $x$ は $0 \leq x \leq 1$ の範囲のみ考えるとします。

DP の遷移を考えます。

頂点 $v$ が葉で、 $U$ または $V$ である場合

頂点 $v$ が葉なので、大小関係に関する制約はありません。

$v = U$ のとき、頂点 $v$ に書かれる数は必ず $a$ なので、頂点 $v$ に書かれる数が $x$ 以下になる確率は $x \lt a$ のとき $0$ , $x \geq a$ のとき $1$ となります。よって、 $\text{dp}[v](a,b,x) = \begin{cases} 0 & (x \lt a) \\ 1 & (x \geq a) \end{cases}$ となります。

同様に、 $v = V$ のときは $\text{dp}[v](a,b,x) = \begin{cases} 0 & (x \lt b) \\ 1 & (x \geq b) \end{cases}$ となります。

頂点 $v$ が葉で、 $U$ でも $V$ でもない場合

上と同じく、大小関係に関する制約はないので、頂点 $v$ に書かれる数は区間 $[0,1]$ から一様ランダムに選ばれます。よって、頂点 $v$ に書かれる数が $x$ 以下になる確率は $x$ となるので、 $\text{dp}[v](a,b,x)=x$ となります。

頂点 $v$ が葉でなく、 $U$ または $V$ である場合

$v = U$ のとき、頂点 $v$ に書かれる数は必ず $a$ です。よって、 $x \lt a$ のとき明らかに $\text{dp}[v](a,b,x)=0$ となります。

$x \geq a$ とします。頂点 $v$ の部分木で大小関係の条件が満たされることは、以下の条件がともに成り立つことと同値になります。

頂点 $2v$ に書かれた数が $a$ 以下であり、かつ頂点 $2v$ の部分木で大小関係の条件が満たされている。
頂点 $2v+1$ に書かれた数が $a$ 以下であり、かつ頂点 $2v+1$ の部分木で大小関係の条件が満たされている。

よって、頂点 $v$ の部分木で大小関係の条件が満たされる確率は $\text{dp}[2v](a,b,a) \cdot \text{dp}[2v+1](a,b,a)$ です。

以上より、遷移は $\text{dp}[v](a,b,x) = \begin{cases} 0 & (x \lt a) \\ \text{dp}[2v](a,b,a) \cdot \text{dp}[2v+1](a,b,a) & (x \geq a) \end{cases}$ となります。

同様に、 $v = V$ の場合の遷移は $\text{dp}[v](a,b,x) = \begin{cases} 0 & (x \lt b) \\ \text{dp}[2v](a,b,b) \cdot \text{dp}[2v+1](a,b,b) & (x \geq b) \end{cases}$ となります。

頂点 $v$ が葉でなく、 $U$ でも $V$ でもない場合

頂点 $v$ に書く数を $t$ とします。頂点 $v$ の部分木で大小関係の条件が満たされることは、以下の条件がともに成り立つことと同値になります。

頂点 $2v$ に書かれた数が $t$ 以下であり、かつ頂点 $2v$ の部分木で大小関係の条件が満たされている。
頂点 $2v+1$ に書かれた数が $t$ 以下であり、かつ頂点 $2v+1$ の部分木で大小関係の条件が満たされている。

これが成り立つ確率は $\text{dp}[2v](a,b,t) \cdot \text{dp}[2v+1](a,b,t)$ です。

$t$ は区間 $[0,1]$ から一様ランダムに選ばれますが、 $t \leq x$ である必要があるので、遷移は $\displaystyle\text{dp}[v](a,b,x) = \int_0^x \text{dp}[2v](a,b,t) \cdot \text{dp}[2v+1](a,b,t) \ dt$ となります。

$\text{dp}[1](a,b,1)$ は、木 $T$ の頂点 $U$ に $a$ を、頂点 $V$ に頂点 $b$ を、他の各頂点に区間 $[0,1]$ に含まれる実数を独立に一様ランダムに選んで書くとき、各辺の大小関係の条件が満たされる確率となります。よって、 $\displaystyle p_2(N, A, B) = \int_0^1 \int_0^a \text{dp}[1](a,b,1) \ db da$ となります。

[4] 計算例

例として、サンプル $3$ の結果を計算する過程を示します。

$\displaystyle\frac{1}{p_1(4)} = 3^4 \cdot 7^2 \cdot 15 = 59535$ である。
頂点 $8$ は葉で頂点 $U$ であるので、 $\text{dp}[8](a,b,x) = \begin{cases} 1 & (x \geq a) \\ 0 & (x \lt a) \end{cases}$ である。以後、このような場合に $\text{dp}[8](a,b,x) = 1 \ (x \geq a)$ と省略する。
頂点 $9$ は葉で頂点 $U$ でも頂点 $V$ でもないので、 $\text{dp}[9](a,b,x)=x$ である。頂点 $10, 11, 12, 13, 14, 15$ も同様である。
$\displaystyle\text{dp}[4](a,b,x) = \int_0^x \text{dp}[8](a,b,t) \cdot \text{dp}[9](a,b,t) \ dt = \int_a^x 1 \cdot t \ dt = \frac{1}{2}\left(x^2-a^2\right) \ (x \geq a)$ である。
$\displaystyle\text{dp}[5](a,b,x) = \int_0^x \text{dp}[10](a,b,t) \cdot \text{dp}[11](a,b,t) \ dt = \int_0^x t \cdot t \ dt = \frac{1}{3}x^3$ である。頂点 $6$ も同様である。
$\text{dp}[7](a,b,x) = \text{dp}[14](a,b,b) \cdot \text{dp}[15](a,b,b) = b^2 \ (x \geq b)$ である。
$\displaystyle\text{dp}[2](a,b,x) = \int_0^x \text{dp}[4](a,b,t) \cdot \text{dp}[5](a,b,t) \ dt = \int_a^x \frac{1}{2}\left(t^2-a^2\right) \cdot \frac{1}{3}t^3 \ dt = \frac{1}{72}\left(2x^6 - 3x^4a^2+ a^6\right) \ (x \geq a)$ である。
$\displaystyle\text{dp}[3](a,b,x) = \int_0^x \text{dp}[6](a,b,t) \cdot \text{dp}[7](a,b,t) \ dt = \int_b^x \frac{1}{3}t^3 \cdot b^2 \ dt = \frac{1}{12}\left(b^2x^4 - b^6\right) \ (x \geq b)$ である。
$\displaystyle\text{dp}[1](a,b,1) = \int_a^1 \frac{1}{72}\left(2t^6 - 3t^4a^2 + a^6\right) \cdot \frac{1}{12}\left(b^2t^4 - b^6\right) \ dt$ である。
$\displaystyle p_2(N,A,B) = \int_0^1 \int_0^a \int_a^1 \frac{1}{72}\left(2t^6 - 3t^4a^2 + a^6\right) \cdot \frac{1}{12}\left(b^2t^4 - b^6\right) \ dt db da$ である。これを計算すると $\displaystyle\frac{89}{38102400}$ となる。
答えは $\displaystyle\frac{89}{38102400} \cdot 59535 = \frac{89}{640}$ である。

[5] 計算量の削減

頂点は $2^N-1$ 個あるため、すべての頂点について $\text{dp}$ を計算することはできません。しかし、木 $T$ には構造が同じ部分が多数存在することを利用すると、 $\text{dp}$ を計算する必要のある頂点を $O(N)$ 個に減らすことができます。各頂点の $\text{dp}$ は $a, b, x$ の多項式となりますが、それらの特徴と計算手順は以下のようになります。ただし、 $P = 998244353$ とします。

$v$ を $U$ の先祖でも $V$ の先祖でもない頂点とすると、 $\text{dp}[v]$ は頂点 $v$ の高さのみに依存する $x$ のみの単項式である。 $\text{dp}[v]$ は $O(N \log P)$ 時間の前計算のもと、 $O(1)$ 時間で得ることができる。
$v$ を $U$ の先祖で $V$ の先祖でない頂点とする。また、頂点 $v, U$ の距離を $d$ とおくと、 $\text{dp}[v]$ は $a, x$ の $d+1$ 項の多項式であり、次数はすべて等しい。 $d>0$ のとき、 $v$ の子で部分木に $U$ を含むほうを $w$ とおくと、 $\text{dp}[v]$ は $\text{dp}[w]$ から $O(d \log P)$ 時間で計算できる。
$v$ を $V$ の先祖で $U$ の先祖でない頂点とする。また、頂点 $v, V$ の距離を $d$ とおくと、 $\text{dp}[v]$ は $b, x$ の $d+1$ 項の多項式であり、次数はすべて等しい。 $d>0$ のとき、 $v$ の子で部分木に $V$ を含むほうを $w$ とおくと、 $\text{dp}[v]$ は $\text{dp}[w]$ から $O(d \log P)$ 時間で計算できる。

これらを用いると、 $\text{dp}[2], \text{dp}[3]$ を $O(N^2 \log P)$ 時間で計算することができます。また、 $\text{dp}[2]$ は $a, x$ の $O(N)$ 項の多項式、 $\text{dp}[3]$ は $b, x$ の $O(N)$ 項の多項式なので、 $\text{dp}[2](a,b,x) \cdot \text{dp}[3](a,b,x)$ は $O(N^2)$ 項の多項式となります。

$\displaystyle p_2(N,A,B) = \int_0^1 \int_0^a \int_a^1 \text{dp}[2](a,b,x) \cdot \text{dp}[3](a,b,x) \ dx db da$ なので、積分を項ごとに分解すると、 $\displaystyle \int_0^1 \int_0^a \int_a^1 x^{e_1}a^{e_2}b^{e_3} \ dx db da$ を計算することに帰着されます。

$0 \leq a \leq 1, 0 \leq b \leq a, a \leq x \leq 1 \Leftrightarrow 0 \leq x \leq 1, 0 \leq a \leq x, 0 \leq b \leq a$ なので、積分の順序を交換して
$\displaystyle \int_0^1 \int_0^a \int_a^1 x^{e_1}a^{e_2}b^{e_3} \ dx db da = \int_0^1 \int_0^x \int_0^a x^{e_1}a^{e_2}b^{e_3} \ db da dx = \frac{1}{(e_3+1)(e_2+e_3+1)(e_1+e_2+e_3+1)}$ となります。これで $p_2(N,A,B)$ を $O(N^2 \log P)$ 時間で求めることができたので、答えも $O(N^2 \log P)$ 時間で求めることができます。

実装

atcoder.jp

#include <bits/stdc++.h>
using namespace std;
const long long MOD = 998244353;
long long modpow(long long a, long long b){
  long long ans = 1;
  while (b > 0){
    if (b % 2 == 1){
      ans *= a;
      ans %= MOD;
    }
    a *= a;
    a %= MOD;
    b /= 2;
  }
  return ans;
}
long long modinv(long long a){
  a %= MOD;
  return modpow(a, MOD - 2);
}
vector<long long> POW2;
struct poly{
  vector<tuple<long long, long long, long long>> P;
  poly(): P(){
  }
  int size(){
    return P.size();
  }
  void multiply_monomial(poly f){
    int N = P.size();
    for (int i = 0; i < N; i++){
      get<0>(P[i]) *= get<0>(f.P[0]);
      get<0>(P[i]) %= MOD;
      get<1>(P[i]) += get<1>(f.P[0]);
      get<1>(P[i]) %= MOD;
      get<2>(P[i]) += get<2>(f.P[0]);
      get<2>(P[i]) %= MOD;
    }
  }
  void integrate(){
    int N = P.size();
    for (int i = 0; i < N; i++){
      get<1>(P[i])++;
      get<1>(P[i]) %= MOD;
      get<0>(P[i]) *= modinv(get<1>(P[i]));
      get<0>(P[i]) %= MOD;
    }
  }
  void integrate_a(){
    int N = P.size();
    long long sum = 0;
    long long deg = (get<1>(P[0]) + get<2>(P[0]) + 1) % MOD;
    for (int i = 0; i < N; i++){
      get<1>(P[i])++;
      get<1>(P[i]) %= MOD;
      get<0>(P[i]) *= modinv(get<1>(P[i]));
      get<0>(P[i]) %= MOD;
      sum += MOD - get<0>(P[i]);
      sum %= MOD;
    }
    P.push_back(make_tuple(sum, 0, deg));
  }
};
int main(){
  int N, A, B;
  cin >> N >> A >> B;
  POW2.resize(N + 1);
  POW2[0] = 1;
  for (int i = 0; i < N; i++){
    POW2[i + 1] = POW2[i] * 2 % MOD;
  }
  vector<poly> f(N);
  f[N - 1].P.push_back(make_tuple(1, 1, 0));
  for (int i = N - 2; i >= 0; i--){
    f[i] = f[i + 1];
    f[i].multiply_monomial(f[i]);
    f[i].integrate();
  }
  poly ga;
  if (A == N - 1){
    ga.P.push_back(make_tuple(1, 0, 0));
  } else {
    ga = f[A + 1];
    ga.multiply_monomial(ga);
    swap(get<1>(ga.P[0]), get<2>(ga.P[0]));
  }
  for (int i = A; i >= 2; i--){
    ga.multiply_monomial(f[i]);
    ga.integrate_a();
  }
  poly gb;
  if (B == N - 1){
    gb.P.push_back(make_tuple(1, 0, 0));
  } else {
    gb = f[B + 1];
    gb.multiply_monomial(gb);
    swap(get<1>(gb.P[0]), get<2>(gb.P[0]));
  }
  for (int i = B; i >= 2; i--){
    gb.multiply_monomial(f[i]);
    gb.integrate_a();
  }
  int N1 = ga.size();
  int N2 = gb.size();
  long long p2 = 0;
  vector<long long> T(N2);
  for (int i = 0; i < N2; i++){
    T[i] = modinv(get<2>(gb.P[i]) + 1);
  }
  for (int i = 0; i < N1; i++){
    for (int j = 0; j < N2; j++){
      long long c = get<0>(ga.P[i]) * get<0>(gb.P[j]) % MOD;
      long long x = (get<1>(ga.P[i]) + get<1>(gb.P[j])) % MOD;
      long long a = get<2>(ga.P[i]);
      long long b = get<2>(gb.P[j]);
      p2 += T[j] * modinv(a + b + 2) % MOD * modinv(x + a + b + 3) % MOD * c % MOD;
      p2 %= MOD;
    }
  }
  long long p1inv = 1;
  for (int i = 2; i <= N; i++){
    p1inv *= p1inv;
    p1inv %= MOD;
    p1inv *= POW2[i] - 1;
    p1inv %= MOD;
  }
  cout << p2 * p1inv % MOD << endl;
}

解法

[1] 条件の言い換え

[2] の計算

[3] の計算

頂点 が葉で、 または である場合

頂点 が葉で、 でも でもない場合

頂点 が葉でなく、 または である場合

頂点 が葉でなく、 でも でもない場合