ABC220E - Distance on Large Perfect Binary Tree

考え方

根を「深さ$0$」とすると,深さ$d$にある頂点の個数は$2^{d}$個.深さ$0$から深さ$d$までの頂点の総数は
\begin{aligned}
\sum_{i = 0}^{d} 2^{i} = 2^{d + 1} - 1
\end{aligned}
であるから,問題文の設定では深さは$N - 1$である.

「①一番浅い頂点$v$」に対して,「②$v$の左側の選び方」と「③$v$の左側の選び方」を数え上げれば良い.

これは,

  1. 左側の深さ$d\in \{0,1,2,...,D\}$を決め,
  2. 左側の頂点の選び方・右側の頂点の選び方を求め,
  3. 深さの制限から可能な頂点$v$の数を求める
ことで実現できる.

最後の可能な頂点数は,一番深い深さ$d$がわかれば,

\begin{aligned}
\sum_{i = 0}^{d} 2^{i} = 2^{d + 1} - 1
\end{aligned}
となる.

【参考】Submission #26145759 - AtCoder Beginner Contest 220

回答例

N, D = map(int, input().split())
M = 998244353

memo = [pow(2, i, M) for i in range(10 ** 6)]

cnt = 0
for l in range(D + 1):
  if max(l, D - l) > N:
    continue
  numl = memo[l - 1] if l > 0 else 1
  numr = memo[D - l - 1] if l < D else 1
  maxd = N - 1 - max(l, D - l)
  cnt += (memo[maxd + 1] - 1) * numl * numr
  cnt %= M
    
cnt *= 2
cnt %= M
print(cnt)