Day 100-Q1 | モンテカルロ木探索の基礎 (UCB1)

問題

モンテカルロ木探索（MCTS）の中核をなす UCB1選択則（多腕バンディット問題の代表的解法）を、厳密なシミュレーションとして実装せよ。

$K$ 本の腕（スロットマシン）があり、腕 $i$ を引くと $0$ または $1$ の報酬が得られる。合計 $T$ 回、腕を選んで引ける。各腕 $i$ の報酬はあらかじめ決められた長さ $T$ の系列 $r_{i,0},\dots,r_{i,T-1}$ として与えられ、腕 $i$ を $c$ 回目に引いたとき得られる報酬は $r_{i,c}$ である（毎回同じ系列を順番に消費する）。

方針: 最初の $K$ 回は腕 $1,\dots,K$ を番号順に1回ずつ引く（ウォームスタート）。以降、そのときまでの総試行回数を $t$ として

$$\text{score}_i = \bar{x}_i + \sqrt{\frac{2\ln t}{n_i}}$$

が最大の腕を引く（同点は番号が小さい方）。$T$ 回終了後の獲得報酬合計を出力せよ。

入力形式

K T
r_{1,0} ... r_{1,T-1}
:
r_{K,0} ... r_{K,T-1}

制約

$2 \le K \le 10$

$K \le T \le 10^5$

$r_{i,c} \in \{0, 1\}$

すべて整数入力

入出力例

入力例1

2 5
1 0 1 0 0
0 1 1 1 0

出力例1

ウォームスタートで腕1(報酬1),腕2(報酬0)。以降UCB1スコアに従い腕1,腕1,腕2の順に選ばれ、合計報酬は $1+1+0+1+0=3$。

概念図

ヒント（段階的開示）

ヒント1: 方向性

平均報酬が良い腕だけを引き続ける貪欲法だと、序盤にハズレを引いた本当は良い腕を二度と試さなくなる危険がある。逆に毎回ランダムでは学習した情報を活かせない。「探索」と「活用」の両立を考えよ。

ヒント2: アプローチ

UCB1は各腕のスコアを「平均報酬（活用）」+「試行回数が少ないほど大きくなる不確実性ボーナス（探索）」の和で表す。$\sqrt{2\ln t/n_i}$ は $n_i$ が小さいほど、$t$ が増えるほど大きくなり、あまり引いていない腕を定期的に見直させる。理論的に後悔（regret）が $O(\log T)$ に抑えられることが知られる。

ヒント3: 誘導（コード骨格）

best_i, best_score = -1, -1.0
for i in range(K):
    mean = s[i] / n[i]
    score = mean + math.sqrt(2 * math.log(t) / n[i])
    if score > best_score + 1e-15:
        best_score = score
        best_i = i

t は「これから引く前の、すでに完了した試行回数」を使う（ウォームスタート直後は t == K）。

模範解答 (Python)

import sys
import math


def solve():
    data = sys.stdin.read().split()
    idx = 0
    K = int(data[idx]); idx += 1
    T = int(data[idx]); idx += 1
    rewards = []
    for _ in range(K):
        rewards.append([int(x) for x in data[idx:idx+T]])
        idx += T

    n = [0] * K
    s = [0] * K

    for i in range(K):
        r = rewards[i][n[i]]
        s[i] += r
        n[i] += 1

    t = K
    while t < T:
        best_i, best_score = -1, -1.0
        for i in range(K):
            mean = s[i] / n[i]
            score = mean + math.sqrt(2.0 * math.log(t) / n[i])
            if score > best_score + 1e-15:
                best_score = score
                best_i = i
        r = rewards[best_i][n[best_i]]
        s[best_i] += r
        n[best_i] += 1
        t += 1

    print(sum(s))


solve()

計算量: $O(KT)$（$K\le10$ なので実質 $O(T)$）。メモリ $O(KT)$。

Step-by-Step 解説

1入力の一括読み込み
$K\times T$の報酬テーブルを一括取得し、行ごとのinput()呼び出しコストを避ける。

2ウォームスタート
全腕を番号順に1回ずつ引き、$n_i \ge 1$を保証してゼロ除算を防ぐ。

3UCB1スコア計算と選択
毎回 $O(K)$ で全腕を走査し最大スコアの腕を選ぶ（同点は番号が小さい方）。

4報酬の消費と統計更新
選ばれた腕の系列から「まだ使っていない先頭」を消費し、$n_i,s_i$を更新。

5出力
$T$回終了後の総獲得報酬 $\sum_i s_i$ を出力する。

よくあるミス

ミス	原因	正しい書き方
`t`に1-indexedの回数を使う	ウォームスタート直後は`t=K+1`から始めてしまう	ループ開始時`t=K`とし、1回引くごとに`t+=1`
同点時のタイブレークを実装しない	浮動小数の丸め誤差でスコアが前後する	`score > best_score + 1e-15`として先着（番号が小さい方）を優先
報酬系列を使い切ってから同じ腕を引く実装ミス	`n[i]`の範囲チェック漏れ	制約上`r_i`の長さは常に`T`保証、自作テストで境界を確認
`math.log(t)`に`t=0`を渡す	ウォームスタート前や初期化ミス	ウォームスタート後は`t>=K>=2`を確認

次のステップ

発展: 報酬分布が途中で変化する非定常バンディットに対応させる（減衰平均を使うUCB変種）
発展: 実際のゲーム木（各ノードがバンディット）にUCB1を再帰適用する完全なMCTS（Selection→Expansion→Simulation→Backpropagation）に拡張する
次回予告: Boyer-Moore法（文字列探索・Bad Character則+Good Suffix則）

自己評価

完全理解おおむね理解要復習

Day 100-Q1 — モンテカルロ木探索の基礎: UCB1多腕バンディット

問題