feature image

2026年4月21日 | ブログ記事

木上のビームサーチ:高速化編

先週、 AHC063 で優勝しました。 AHC 初優勝。しかも長期。感無量です。

ところで、今回の問題に取り組んでいる過程で次のような状況が発生しました。
多くのテストケースでおそらく最適解っぽいやつ、もしくは最適解の付近の解が出ている。さらに加えて、評価が順位スコアで、上位争いをするとなると難しいケースで差をつけても、簡単なケースでできうる限り失敗することなく安定して解けなければ話にならない。

このような状況でのもっともインスタントな目標は高速化です。どーせ最適解付近の解なんて人間がうまいパターンを見いだせるわけないし、賢い構築とか考えるだけ無駄、速度でなぐってしまうのが堅実で確実という判断ですね。残念なことに、純粋な定数倍高速化というのはすぐにサチりがちで、探索空間を大きく変えない、定数倍高速化と似たような効果を持つ細かなヒューリスティックの重ねがけ勝負になることが多いかと思います。

ただ、今回の私の解法で中心的な役割を果たしている木上のビームサーチ、これは純粋な定数倍高速化の余地がかなりあるということで有名(?)で、今回の記事は、もしかしたら、今回の優勝に少しは寄与しているかもしれない、そんな木上のビームサーチの定数倍高速化についてです。
詳しい方向けに要約すると、「木上のビームサーチをオイラーツアーで管理するやつ、実は帰りがけ順に頂点を並べたやつで十分だからそういうふうに実装すると速くなるかも!?」という内容になります。

木上のビームサーチ

木上のビームサーチは、ビーム幅の個数分の選んだ頂点の履歴(の木)を気合いで管理して、必要な状態をコピーではなくすべて差分更新で列挙することで、めちゃ高速にビームサーチができる、というやつです。

今回の実装は特に世代のスキップがない場合に特化した実装になります。
世代のスキップがある場合というのは、例えば MP 消費 1,2,3 の 3 種の魔法があり毎ターンこれらの魔法のどれかを使うという問題で、残り MP が同じやつのなかで評価値上位を選ぶという場合ですね。こういうのは今回の実装のサポート対象外です。逆に経過ターンで揃えて上位を選ぶというビームサーチをする場合は可能になります。ターンはどのような行動をしても必ず1ずつ進みますので。

履歴の木の管理が曲者で、オイラーツアーを持つとか、二重連鎖木で管理するとか色々やり方があります。

詳しく説明すると、

わざわざ NodeCand を分けているのは、評価値やハッシュなどは State の差分更新には必要なく、そして Node はたくさんコピーすることになるので最小限のデータだけを持たせて高速化を図っているためですね。
これらの構造体、関数が用意できているときに、

この 3 つの機能を持つ BeamSearch 構造体を作るのを目標にしましょう。

実装

の 5 つの変数を使います。 expand していく過程で次の run に用いる tour,leaf,cand も同時に作るのでそれを next_tour,next_leaf,next_cand で管理します。
イメージ図:

赤い頂点が trace に入っている頂点、黒い頂点が tour に入っている頂点で頂点に書かれている数字はインデックスを表します。
leaf,cand は図の左から右の方向にインデックス昇順になっており、選択されている cand が赤くなっています。
また、 leaf[i+1]-leaf[i] が左から i 番目の葉と i+1 番目の葉の最近共通祖先までの距離に対応しています。(これ、重要です)
注意点として、 tour で頂点が連続する区間内では(通常の帰りがけ順とは異なり)葉に近い方がインデックスが大きくなるように管理します。 tourtrace と同じ向きにすることで Node のコピーが簡略化できます。

cand から次の cand まで移動するにはこのようにします:

青い頂点が next_tour に入っている頂点になります。

以上を実装するとこうなります:

fn expand(state:&mut State,cand:&Cand,leaf:usize,next_cand:&mut Vec<Cand>){
    let mut add_cand=|new:Cand|{
        assert!(new.leaf==leaf);
        assert!(!new.active);
        next_cand.push(new);
    };
    
    todo!();
}


struct State{}
impl State{
    fn apply(&mut self,node:&Node){
        todo!();
    }
    
    fn revert(&mut self,node:&Node){
        todo!();
    }
}


#[derive(Clone)]
struct Cand{
    leaf:usize, // この cand の親の leaf のインデックス
    active:bool, // この cand が次に expand するかを管理
}
impl Cand{
    fn to_node(&self)->Node{
        Node{}
    }
}


#[derive(Clone,Copy)]
struct Node{}


struct BeamSearch{
    state:State,
    trace:Vec<Node>,
    tour:Vec<Node>,
    leaf:Vec<usize>,
    cand:Vec<Cand>,
    next_tour:Vec<Node>,
    next_leaf:Vec<usize>,
    next_cand:Vec<Cand>,
}
impl BeamSearch{
    fn new(state:State,initial:Cand)->Self{
        Self{
            state,
            trace:vec![],
            tour:vec![],
            leaf:vec![0],
            cand:vec![initial],
            next_tour:vec![],
            next_leaf:vec![],
            next_cand:vec![],
        }
    }
    
    fn run(&mut self){
        let turn=self.trace.len();
        self.trace.push(self.cand.last().unwrap().to_node()); // trace の長さを増やすために適当に一つ追加
        
        if turn==0{ // 初回はすでに state が initial の指し示す状態になっているため特別扱い
            expand(&mut self.state,&self.cand[0],0,&mut self.next_cand);
            std::mem::swap(&mut self.cand,&mut self.next_cand);
            return;
        }
        
        self.next_tour.clear();
        self.next_leaf.clear();
        self.next_cand.clear();
        
        let mut f=0; // 最初の revert では trace[turn] を revert してはいけないのでそれを管理
        let mut li=self.leaf.len()-1; // state が leaf[li] まで進んでいることを管理
        
        for cand in self.cand.iter().rev().filter(|cand|cand.active){
            let lca=self.leaf[cand.leaf..li].iter().rev()
                .fold((0,self.leaf[li]),|(a,b),&x|(a.max(b-x),x)).0; // lca までの距離を計算
            self.trace[turn-lca..turn+f].iter().rev().for_each(|node|self.state.revert(node)); // lca まで戻る
            f=1;
            self.next_tour.extend_from_slice(&self.trace[turn-lca..]); // 帰りがけ順を次の tour に追加
            
            self.trace[turn]=cand.to_node();
            let mut prog=0;
            for w in self.leaf[cand.leaf..=li].windows(2){
                let rank=w[1]-w[0];
                if prog<rank{
                    self.trace[turn-rank..turn-prog].copy_from_slice(&self.tour[w[0]..w[1]-prog]); // cand の表す状態までの経路を計算
                    prog=rank;
                }
            }
            
            self.trace[turn-lca..].iter().for_each(|node|self.state.apply(node)); // cand の表す状態まで進める
            expand(&mut self.state,cand,self.next_leaf.len(),&mut self.next_cand);
            self.next_leaf.push(self.next_tour.len());
            li=cand.leaf;
        }
        
        std::mem::swap(&mut self.tour,&mut self.next_tour);
        std::mem::swap(&mut self.leaf,&mut self.next_leaf);
        std::mem::swap(&mut self.cand,&mut self.next_cand);
    }
    
    fn restore(&self,idx:usize)->Vec<Node>{ // cand[idx] までの経路を復元
        let mut ret=self.trace[1..].to_vec();
        
        let len=ret.len();
        let mut prog=0;
        for w in self.leaf[self.cand[idx].leaf..].windows(2){
            let rank=w[1]-w[0];
            if prog<rank{
                ret[len-rank..len-prog].copy_from_slice(&self.tour[w[0]..w[1]-prog]);
                prog=rank;
            }
        }
        
        ret.push(self.cand[idx].to_node());
        ret
    }
    
    fn select(&mut self,i:usize){
        self.cand[i].active=true;
    }
}

なおテストなのですが、コーナーケースとしてビーム幅が 1 の場合や expand で子が一つも追加されない場合などがあります。実装のお供にしてみてください。

おわりに

ちなみにこの実装は AHC063 のコンテスト終了後、木上のビームサーチに感謝を捧げているときに思いつきました。
最終提出のビームサーチ部分をこの記事の実装に置き換えると、ケースにもよりますが 20% 程度の高速化になるようです。
problem : https://atcoder.jp/contests/ahc063/tasks/ahc063_a
before : https://atcoder.jp/contests/ahc063/submissions/74935244
after : https://atcoder.jp/contests/ahc063/submissions/75152063
ライブラリパート以外でもそこそこ時間がかかってるのでかなり偉い気がします。(元の実装があまり良くなかったのもありそうですが...。)

木上のビームサーチ verify 問題であるところのゲーム実況者Xの挑戦にも提出してみました。
problem : https://atcoder.jp/contests/rco-contest-2018-qual/tasks/rco_contest_2018_qual_a
submission : https://atcoder.jp/contests/rco-contest-2018-qual/submissions/75152175
329K の大台を突破。これ可能だったのか...。

ここからさらに高速化するとなると、 Node のコピーを少なくする方向になるのでしょうか?不可能な気がして諦めてしまいました。
あと目下の課題は世代のスキップがある場合の木上のビームサーチですね。ポインタ木が本当に遅くて困ってます。もしかしたら、次の探索に必要のない頂点のコピーが発生するのは諦めてオイラーツアーで管理したほうが速いのでは?良い実装を思いついた人は教えてください。これが遅いせいでスキップありの世代の取り方に抵抗を感じるようになってしまいました。





おわりです。
jibjib icon
この記事を書いた人
jibjib

この記事をシェア

このエントリーをはてなブックマークに追加
共有

関連する記事

2021年4月27日
AHC のおはなし
NNMochi icon NNMochi
2025年8月19日
AHC051 の記録
jibjib icon jibjib
2023年12月10日
HACK TO THE FUTURE 2024 に参加しました!
noya2 icon noya2
記事一覧 タグ一覧 Google アナリティクスについて 特定商取引法に基づく表記