メモリを節約しつつ動的計画法の経路を復元する

背景

動的計画法を実行する際、計算した部分問題の解をテーブルに書き込んでいくわけだが、問題によっては、テーブルの全体を最後まで持っておく必要がないこともある。

例として、2つの文字列のレーベンシュタイン距離を求める問題を考える。二つの文字列の先頭それぞれi文字とj文字を取ったもののレーベンシュタイン距離をdp[i][j]とすると、dp[i][j]の計算にはdp[i-1][j]dp[i-1][j-1]dp[i][j-1]の値があれば十分なので、iが小さい方から計算していくなら、dp[i][j]を計算する時点でdp[i-2][*]dp[i-3][*]などは忘れてしまっていても良い。コードは例えば以下のようになる。

use std::cmp;
use std::iter::Iterator;
use std::mem;

// 二つの配列のレーベンシュタイン距離を求める
fn levenshtein<T: Eq>(x: &[T], y: &[T]) -> usize {
    let n = x.len();
    let m = y.len();
    let mut prev: Vec<usize> = (0..=m).collect(); // 一個前の列
    let mut this: Vec<usize> = vec![0; m + 1]; // これから計算する列

    for i in 1..=n {
        // ここで prev[j] には dp[i-1][j] が入っている
        // this[j] にはこれから dp[i][j] が入る(今のところthisの内容はなんでも良い)
        this[0] = i;
        for j in 1..=m {
            if x[i - 1] == y[j - 1] {
                this[j] = prev[j - 1];
            } else {
                this[j] = 1 + cmp::min(this[j - 1], cmp::min(prev[j], prev[j - 1]));
            }
        }
        mem::swap(&mut prev, &mut this);
    }

    prev[m]
}

これを図示すると次のような絵を描くことができる。

<rdf:RDF> <cc:Work rdf:about=""> <dc:format>image/svg+xml</dc:format> <dc:type rdf:resource="http://purl.org/dc/dcmitype/StillImage" /> <dc:title></dc:title> </cc:Work> </rdf:RDF>

経路復元が必要な場合

ここからが本題。

DAG上の最短・最長距離を求めていると解釈できるタイプの動的計画法では、テーブルの各要素について、そこに至る最適経路(のうち一本)がどこから来ているかを記録しておくことで、最大・最小値だけでなく、それを実現する経路を計算することができる。レーベンシュタイン距離の例なら、距離を求めるだけでなく、その距離を実現する編集操作の列を具体的に求めることができる。

これを素朴に実装すると、テーブルの要素数に比例するメモリが必要になる。経路の復元はテーブルを埋め終わった後にしか行なえないので、経路情報を忘れてしまうことが許されない。

ところが、平方分割のアイディアを使うと、メモリ使用量を大幅に減らすことができる。具体的には、経路復元が必要ない場合に1/nの省メモリが達成できなら、1/sqrt(n)の省メモリを達成しつつ経路復元をすることができる。

これをするには、テーブルを埋めていく際、古い要素を完全に忘れてしまう代わりに、適当な間隔を開けて一部の要素を覚えておき、後でそこからテーブル埋めを再開できるようにする。

<rdf:RDF> <cc:Work rdf:about=""> <dc:format>image/svg+xml</dc:format> <dc:type rdf:resource="http://purl.org/dc/dcmitype/StillImage" /> <dc:title></dc:title> </cc:Work> </rdf:RDF>

テーブル埋めが完了した後は、覚えておいた要素からテーブルを埋め直しながら経路を復元していく。

<rdf:RDF> <cc:Work rdf:about=""> <dc:format>image/svg+xml</dc:format> <dc:type rdf:resource="http://purl.org/dc/dcmitype/StillImage" /> <dc:title></dc:title> </cc:Work> </rdf:RDF>

レーベンシュタイン距離の問題では、二つの文字列の長さをm, nとすると、覚えておく要素の間隔を{\sqrt{n}}にすることで、空間計算量を{\mathcal{O}(m\sqrt{n})}にできる。

コード例は以下。

use std::cmp;
use std::iter::Iterator;

// 文字列に対する編集操作
#[derive(Copy, Clone, Debug)]
enum DiffOp {
    Insert,
    Delete,
    Change,
    Keep,
}

// DPテーブルの内容
#[derive(Copy, Clone, Debug)]
struct Entry {
    op: DiffOp,
    score: usize,
}

// 二つの配列のレーベンシュタイン距離を実現する操作(挿入・削除・置換)
// の列を求める。
fn levenshtein_diff<T: Eq>(x: &[T], y: &[T]) -> Vec<DiffOp> {
    let n = x.len();
    let m = y.len();

    // dp[i-1][*] の配列から dp[i][*] のVecを作る。
    let forward = |i: usize, tbl: &[Entry]| -> Vec<Entry> {
        let mut ret = Vec::new();
        ret.reserve(m + 1);
        ret.push(Entry {
            op: DiffOp::Delete,
            score: i,
        });
        for j in 1..=m {
            if x[i - 1] == y[j - 1] {
                ret.push(Entry {
                    op: DiffOp::Keep,
                    score: tbl[j - 1].score,
                });
            } else {
                let entry = *[
                    Entry {
                        op: DiffOp::Delete,
                        score: tbl[j].score + 1,
                    },
                    Entry {
                        op: DiffOp::Insert,
                        score: ret[j - 1].score + 1,
                    },
                    Entry {
                        op: DiffOp::Change,
                        score: tbl[j - 1].score + 1,
                    },
                ]
                .iter()
                .min_by_key(|e| e.score)
                .unwrap();
                ret.push(entry);
            }
        }
        ret
    };

    // 第一パス。dp[n][m]を目指して計算しつつ、iがchunk_sizeの倍数のときは
    // dp[i][*]を保存しておく。

    let chunk_size = ((n + 1) as f64).sqrt() as usize;
    let mut saved: Vec<Vec<Entry>> = // 保存しておく配列
        vec![
        (0..=m)
        .map(|j| Entry {
            op: DiffOp::Insert,
            score: j,
        })
        .collect()
        ];
    let mut current: Vec<Entry>;
    let mut prev = &saved[0];

    for i in 1..=n {
        current = forward(i, prev);
        if i % chunk_size == 0 {
            saved.push(current);
            prev = saved.last().unwrap();
        } else {
            prev = &current
        }
    }

    // 第二パス。dp[n][m]から経路を辿っていく。DPテーブルの大部分は忘れて
    // しまっているので、保存したところから再計算する。

    let mut rev_path = Vec::new(); // 復元中の経路。逆順で。
    let mut j = m; // 現在復元している位置のj座標

    for (chunk_no, saved_col) in saved.into_iter().enumerate().rev() {
        // 再計算
        let mut chunk_table: Vec<Vec<Entry>> = vec![saved_col]; // 再計算したテーブル
        let i_end = cmp::min((chunk_no + 1) * chunk_size, n + 1);
        for i in chunk_no * chunk_size + 1..i_end {
            chunk_table.push(forward(i, chunk_table.last().unwrap()));
        }
        // 経路復元
        for (col_no, col) in chunk_table.iter().enumerate().rev() {
            loop {
                if (chunk_no, col_no, j) == (0, 0, 0) {
                    // dp[0][0]については、復元すべき経路はない。
                    break;
                }
                rev_path.push(col[j].op);
                match col[j].op {
                    DiffOp::Insert => {
                        j -= 1;
                    }
                    DiffOp::Change | DiffOp::Keep => {
                        j -= 1;
                        break;
                    }
                    DiffOp::Delete => break,
                }
            }
        }
    }

    rev_path.into_iter().rev().collect()
}

なぜこれを書いたか

実用コードを書いていたら出てきたので。それなりに汎用的な方法だし、競技プログラミング業界では広く知られているのではないかと思うが、ちょっとググっても出てこなかったので記事にした。

マージコンフリクトを高速に解決するためのvimプラグイン

多くの場合、gitのマージコンフリクトの解決は簡単な作業である。これに生活の約1%を費やしているのだが、どうも適切な道具がないせいで効率が悪い気がしていた。たとえばvimdiffは、二つのファイルを比較するのであればとても優秀なのだが、マージコンフリクトを素朴に扱うと4つものバージョン(マージしようとしている2つと、その共通祖先、そして現在のファイル)を比較することになり、いたるところがハイライトされて訳が分からなくなったり、行の対応が崩れるといった問題が高頻度で発生する。また、gitのマージ処理は行単位であるが、単語単位や文字単位で差分を見れば実は自動で解消しうるコンフリクトも少なくない。こういうものは自動で処理してほしい。

こういうことをするvimプラグインは当然すでに存在するだろうと思って探したのだが、見つからなかったので自分で書いた

こんな感じで使う。

デモ画像

見て分かるように、特にペインを分けたりせず、コンフリクトマーカのあるファイル上で直接作業する。ただし重要な点が一つあって、マージしようとしている二つのバージョン(以下A、B)だけでなく、その共通祖先(正確にはmerge-base)もコンフリクトに含まれるようにしておく必要がある(git config --global merge.conflictStyle diff3)。これがないと、どんな変更をマージしようとしているか分からないので、このプラグインには何もできない*1

:Conflict3Highlightコマンドで、diffによる色付けをしている。上のデモでは、Aと共通祖先を比較してその差分を青で塗り、次にBと共通祖先を比較してその差分を水色で塗っている。

この色分けを頼りに手動でコンフリクトを解消しても良いし、このデモのように、プラグインが提供するコマンドを使って半自動で解消しても良い。

半自動での解消は、コンフリクトを段階的に単純化・縮小していくことによって行なう。最終的にコンフリクトが空になった時が解消完了である。

:Conflict3ResolveOneコマンドは、コンフリクトから自動解消できる部分を探して、三つのバージョン(A、B、祖先)でその部分が同じになるようにコンフリクトを単純化する。:Conflict3ResolveAllコマンドは、これを可能な限り繰り返す。:Conflict3Shrink!コマンドは、既に解消されてバージョン間の差がなくなった行をコンフリクトの外に出す。これによって空になったコンフリクトは消される。

実装

実装について特筆すべきことは、diffアルゴリズムvim scriptで自前実装している(行単位はA*、文字単位はMyers法)ことくらいだろうか。これによって、diffアルゴリズムの微妙なカスタマイズが可能になっている。例えば、文字単位diffでは、細かいhunkが増えるのを防ぐために、hunk一個ごとにペナルティを課している。

しかし実装作業としては、diffアルゴリズムよりも、ハイライトや編集作業のこまごまごした処理のために大量の場合分けを書く方が苦しかった。

今のところコードの質はかなりひどくて、クソの山と言っても良いと思う。これはいずれ改善する。

結論

ハイライト部分だけは数ヶ月前に完成していたので、ずっと実用していたのだが、既に手放せなくなっている。とても便利なプラグインになったと思うので、gitを扱うvim使いにはぜひ試してほしい。

vimプラグインを書くのは初めてなので、なにか変なことをしていたら教えていただけるとありがたい。

*1:このプラグインを使わない場合でもこれを指定しない理由はないように思う

ディレクトリを開くコマンドを拒否する

最近のvimは、:editや:splitなどのコマンドでディレクトリを開こうとするとnetrwを起動する。しかし、私がこれをやるのは九割九分、tab補完の後すぐにEnterを叩いたものの補完されたものが予想と違った場合で、ディレクトリを開くのは邪魔な挙動でしかないのでやめさせたいと思った。ただnetrwを無効化するだけではだめで、ディレクトリを開こうとした時にはを完全に無視してコマンドモードに留まって欲しい。

ググっても解決法が見つからない*1ままになっていたが、ちょっとした工夫で実装できることに気がついた。

cmap <CR> <SID>(setcr)<SID>(cr)
cnoremap <SID>(setcr) <C-\>e<SID>Setcr()<CR>

function! s:Setcr()
  let cmdline = getcmdline()
  if match(cmdline, '\v^(e|edit|r|read|sp|split|vsp|vsplit)>.*/$') != -1
    cnoremap <SID>(cr) <NOP>
  else
    cnoremap <SID>(cr) <CR>
  endif
  return cmdline
endfunction

マッピングの中でマッピングを変更するというややこしいことになっているが、これより簡単な方法は思い付いていない。

*1:検索語の選択も難しい

ICFP Programming Contest 2015

今年もICFP Programming Contestに参加していた。今回は「Japanese Curry.」というチームで、メンバーはA氏とJ氏と私の三人だった。

今回の問題テトリス風のゲームをプレイするAIを作るというもの。以下のような特徴がある。

  • ゲーム盤が六角形の敷き詰めで構成されている。
  • 盤の大きさや落ちてくるピースの形が一つに決まっておらず、パラメタになっている。これらは問題ごとに指定される。
  • 次にどのピースがやってくるかを選ぶ疑似乱数が決まっているので、ゲーム終了まで先読みできる。決められた数のピースを消化するか、死ぬと終了。
  • AIからゲームへの入力をコマンド列として与えるのだが、それを文字列として解釈したときに特定の句(phrases of power。コトダマ?)が含まれていると回数に応じて点が得られる。これは主に(全部?)ラヴクラフトの作品からの引用になっているが、問題文中には一つしか明記されておらず、主催者のヒントを頼りに探すことになる。

例年に比べると単純な問題だったのでモチベーションを維持できるか心配だったが、終わってみるといろいろ工夫を要求されて楽しかった。クトゥルフ要素も良いフレーバーになっていた。

やったことは以下。

+0:00 問題文を読む。テトリスだ。

+1:00 入力パーサを書く。といってもJSONなのでaeson一発。ここで既存フォーマットを使うのは核心部分以外に費やされる実装労力を極力削減しようとしているようで、去年までの問題の傾向との違いを感じる。就寝。

+11:00 起きたらA氏がUnicode文字の六角形⬢⬡とANSIカラーシーケンスを組み合わせた美しい盤面表示関数を完成させていた。これは単なる文字列なのでトレース出力などに埋め込むことができ、最後までデバッグに貢献した。続いてA氏は斜交座標を併用するべきだと主張し、J氏とともに実装を開始した(これは終わってみるとまったくその通りで、問題文にある座標系だけを使っていたら手間が爆発してしまう)。あまりやることがないので、提出用コードのコマンドライン解析部分を書き、ひたすらピースを下に落とし続ける馬鹿AIを書いて提出してみる。ゲームが終了するのと同時にコマンド列も終了していなければならないらしく、この方法では1点も取れないことが分かる。

+12:00 このあたりでJ氏が離脱。以降二人チームになる。方針を話し合った結果、各ピースについて置ける場所を全列挙し、置き方についてはヒューリスティクス系の探索で選ぶことになった。まず可能な置き方を列挙する関数を書く。深さ優先探索をするだけだったが、ピボットの座標が負になれるなどややこしい要素が多くて時間がかかる。終了状態から初期値へと逆方向に探索すれば、盤面がほぼ空の状態だと主要部分の計算量をO(width * height)からO(width + height)に落とせそうなことに実装後に気付いたが、これを試す機会はなかった。この間にA氏はゲームルールを実装し、最初の馬鹿AIを改良して0点でない点が取れるAIを書いていた。その後やることがなくなったとのことなのでペアプログラミングに移行。

+13:30 回転を使わない置き方については全列挙ができるようになる。ここでA氏は回転の実装を続け、私はビームサーチの実装に取りかかる。

+15:30 ビームサーチの実装が終わり(ただしバグっていて貪欲法の劣化版になっていた)、解を出力できるようになる。しかし、出力を見ても何をやっているか分からないのでステップ実行のできるシミュレータを書く。

+16:30 ルール部分のバグを潰しつつシミュレータを書き終わる。

+18:00 評価関数として、現在のスコア、「置けるピースをなるべく下に置くことを好む」の二つを実装する。これにより、問題0を死なずに最後まで解けるAIが完成する。しかし非常に遅いので時間を取って最適化することになった。

+19:30 いくつかボトルネックを潰すことに成功する。ビームサーチの実装が壊れていたことに気付いて直す。まだコードが遅いので、高速化の方法を考えた結果、盤面の表現を、疎な永続ビット集合であるIntSetから密なビットマップに変更することに決める。その間、A氏は手当り次第に評価関数用の特徴量を実装することになった。

+23:00 A氏が穴の数を数える特徴量を実装した。私のほうはビットマップの実装でバグを量産して一向に実装が進まないので、あきらめてLightning Division用の解を生成しはじめる。この時点ではじめて問題0以外の解を目視で確認する。AIが遅いので、問題によってはビーム幅を1に絞ったり、穴の数を数える特徴量を外したりした。

+24:00 Lightning Division終了。どうやら「穴の数」に重きをおきすぎていることが判明する。穴がなくても、ブロックの積み方が入り組んだ形になるとどうしようもなくなるケースを確認する。就寝。

+31:00 gitのコミットログによると、このころA氏が「表面積」「空きマスの上にあるブロックの数」「上端から始めて、下降と左右移動だけで到達できるマスの数」の三つの特徴量を実装したらしい。

+33:00 引き続き苦しみながらビットマップを実装する。盤面を偶数段と奇数段に分割したあと、それぞれを8x8の小ブロックに分割し、各ブロックを一個の64ビット整数で表現するという方針だったが、実装がひたすら苦しかった。1次元の可変長ビットベクトルを介して実装する方が良かったのだろうか。

+40:00 ついにビットマップの実装が終わり、表面積をビット演算で計算できるようになるが、高速化は2倍にも満たなかった。A氏はビームサーチを見限って新しい探索技法を完成させ、「遺伝的探索」と名付けていた。ビームサーチが単一の評価関数を使うのに対し、「遺伝的探索」は特徴量に対して数十個の異なる重み付けを持ち、それぞれの重みに従って複数の最良候補を選ぶことで解の多様性を確保するというアイディアらしい。探索終了後には遺伝的アルゴリズムに基いて新しい重み付けの集合が提案される。説明を聞いても半信半疑だったが、ビームサーチよりも結果が良いようなので以降これを使うことにする。

+42:30 ピースの置き方を一部重複して生成していた問題を修正したところ、AIが2倍早くなる。

+44:00 表面積の計算がまだ重いので、盤面更新にあわせて差分を計算できるようにする。

+46:00 AIの動きを見ると、必要もないのに穴を塞ぐケースが多いようだった。穴の数を数える特徴量があれば防げそうだが、計算があまりに重いので無効になっていた。何か方法はないか首をひねっていると、A氏が突然、オイラー標数を使えないかと言いだす。聞いてみると、

  • 連結な平面グラフのオイラー標数(頂点数-辺数+面数)は2である
  • なので、盤面のうち空白の頂点と、それらの隣接関係からなるグラフのオイラー標数を数えれば連結成分の数が分かる
  • 頂点数と辺数はビット演算で数えられる(差分計算もできる)。面数のうち単位三角形の数はやはりビット演算や差分計算で数えられる。単位三角形でないものの数は、孤立したブロックの集合(島と呼ぶ)の数に等しい
  • 島は永続Union-findで管理すれば差分計算で数が求まる!

というアイディアだった。発想に驚嘆しつつ早速実装にとりかかる

+49:00 ひととおり実装したもののUnion-findのバグがとれない。A氏は遺伝的アルゴリズムを自動化しつつ重み行列の更新を続けており、得点がかなり向上していた。明示的に教えていないにもかかわらずAIが複数行の同時消しを行なえるようになったのに驚く。就寝。

+57:00 ついにオイラー標数を使った穴の数の計算が動いたので、満を持してphrases of powerを埋め込むコードを書き始める。方針としては、遺伝的探索で発見したピースの置き場所は変えないことにし、各ピースを置く手順を最適化することにする。各段で右か左の多くとも一方にしか動かない、という制約を入れると、phrases of powerを(場合によっては複数並列で)認識して得点を出力する有限状態機械を考えることで、動的計画法によって最適手順が計算できるので、これを実装することにした。この方法だとphrase of powerを一種類使うたびにもらえる追加点を最適化できないが、これをうまく扱うのは後まわしにすることにした。これは結局最後まで実装することはなかった。

+69:00 ほとんど丸一日をかけた動的計画法の実装が終わる。例によって遅いコードで、"ei!"という三文字のphrase of powerだけを認識する状態でも問題0に対して数秒必要だった。高速化を試みたが、大したことはできなかった。

+70:00 時間管理を実装する。AIのパラメタ1から始めて倍々にし、時間が切れた時点で最善のものを出力する方法を実装した。具体的には、遺伝的探索で使う重み付けの数と、動的計画法で考慮するphrases of powerの長さの合計をパラメタで制御できるようにした。この間A氏はphrases of powerを探していたが、予選問題のマップに埋め込まれていた三つと、長大すぎて計算コストの重い51文字のもの以外は発見できずにいた。

+71:00 A氏が提出用の解を生成している間、phrases of powerを求めてラヴクラフトの引用集、Wikipediaクトゥルフ関係のページやThe H.P. Lovecraft Wikiをさまよった。結局発見できたphraseは一個だけだった。

+72:00 終了。提出が間に合わなかった問題がいくつかあり、予選順位は66位だった。問題17と問題19では一位だったが、これについては特別な何かを実装したわけではないので、運が良かっただけなのではないかと思う。

感想

速いコードをバグらせずに書けるようになるべきだと思った。

AtCoderでHaskell解がスタックオーバーフローを起こす問題の回避策

これを書いている時点でAtCoderHaskell処理系はGHC 7.4であり、これはプログラムに対してデフォルトで8MBのスタック制限を課す。したがって、深さ数十万の再帰を普通に書くと容易にスタック溢れによるREを起こす。

このスタック制限値はコンパイル時または実行時にRTSの-Kオプションで変更できるが、AtCoderではGHCにも自分のプログラムにも好きなオプションを渡すことが基本的にはできない*1ので、他の方法を考える。

回避策

以下のコードをコピペした上で、mainの先頭でunlimitStackSizeを呼べば良い。

import Control.Monad
import Data.Word
import Foreign.Ptr
import Foreign.Storable

-- | Remove the stack size limit in the RTS (+RTS -K).
unlimitStackSize :: IO ()
unlimitStackSize = do
  current <- peek maxStackSizePtr
  when (current == 0x100000) $ -- default limit: 1M words
    poke maxStackSizePtr (-1)

maxStackSizePtr :: Ptr Word32
maxStackSizePtr = plusPtr rtsFlagsPtr 12

foreign import ccall "&RtsFlags" rtsFlagsPtr :: Ptr Word32

RTSがオプションの値を保存するのに使っているグローバル変数の値を書き換えているだけだが、構造体内のメンバのオフセットなどをハードコードしているのでバージョンが変わると動作しなくなるかもしれない。もっとも、GHC 7.8でデフォルトのスタック制限が「物理メモリ量の80%」に変更されるので、おそらくこのようなハックは必要なくなるだろう。

(上記のコードは、この日記の他の部分と同様、コピー自由かつ無保証である。念のため)

*1:Template Haskellを使ってコンパイル時にGHCインスタンスをもう一つ起動するなどすれば可能かもしれない

ICFP Programming Contest 2014

ICFPコンテストに参加していた。今年はLightning Divisionのみで、一人チームでの参加。

Lightning Divisionの課題は、あるゲーム*1をプレイするAIを書くというもので、最初はよくある問題のように思ったのだが、結果的にはかなり楽しめた。というのも、提出するAIは決められた言語で書かなければならないのだが、その言語が機械語のようなものであって、複雑なコードを書こうと思うとまずコンパイラを書くことになる。このコンパイラのソース言語はもちろん自由なのだが、機械語の設計からして露骨にLISP系言語が優遇されている。したがって、数時間ででっちあげた自家製の糞LISP方言*2を使って、ひーひー言いながらAIを書くことになる。これがなかなか面白かった。

時系列

+0時間
問題文を読む。
+1時間
問題文を読み終わって方針を考える。一人チームなのでシミュレータは自作せずに公式のものを使い、コンパイラだけ書くことにする。制御構造としてgotoとwhileを持った手続き型言語を書くことも検討したが、まともな手続き型言語の処理系を書いた経験がないので見送り、誘導された通りにLISP風言語を作ることにする。
+3.5時間
とりあえず動くコンパイラができる。コード生成の面倒な部分(レジスタ割り付け、クロージャ変換、タグビットなどなど)は全てCPUがやってくれる仕様なのであっさりできた。パーサは既存プロジェクトの物を流用。これのおかげでインデント構文を持つLISPにり、モチベーションの向上に寄与した。就寝。
+9.5時間
二度寝
+13時間ごろ
AI作成にとりかかる。方針としては、αβ法みたいに数ステップ先読みして、敵やアイテムからの距離をもとに局面評価すれば、のようなことを漠然と考えていた。いずれにしても最短経路探索は必要だろうと考えそれを当面の目標にする。機械語の仕様上配列が使えないが、リストを毎回書き換えるのはさすがに遅いと思われたのでconsセルを使って完全二分木を作ることにした。
+14時間ごろ
公式のJavaScriptシミュレータ上で完全二分木のコードが動いているらしいことを確認。小さい関数を一つ書くたびにシミュレータにコピペしてテストしていたので時間が掛かる。

ここからはひたすら実装で、時間の記憶があやふやになっている。skew heap*3、キュー、幅優先探索を実装し、その過程で必要になったlength,mapといった標準関数を実装した。コンパイラにも随時手を加え、トレース命令やlist,andといった特殊形式を実装した。

+20時間ごろ
AIの最初のバージョンができる。公式シミュレータで動かしてみたところいつまでたっても初期化が終わらない。初期化時に全点対間の最短経路を計算しているのが重いらしい。計算量を概算したところ制限には間に合っているようにしか思えなかったので、単にJavaScript製であろうシミュレータが遅すぎるのだと信じることにする。以降テストは3x4などの小さいマップで行なった。
+21時間ごろ
AIがクラッシュするのをデバッグするには、ゲームのシミュレータよりもCPUシミュレータを使った法がエラーメッセージが詳しいので、AIをCPUシミュレータで動く形式にコンパイルできるように手を加える。
+23時間ごろ
この時点でのAIの挙動は、最も近い敵からの距離を稼ぎつつ最も近い錠剤を目指すもの。ここから一時間で実装できる改善として、果物の出現位置を目指す項を評価関数に入れるか、敵から逃げるにあたって敵が方向転換できない事実を利用できるようにするかを考えた。迷った挙句後者を選択。
+23.9時間ごろ
結局間に合わなかったので古い版を提出。敵の方向を理解するものができたのは締め切りを90秒ほど過ぎてからだった。

まとめ

コンパイラが動くまでは良いペースだったが、その後のAI実装が遅すぎた。提出したAIは一手読みしかしておらず、未来予測らしいことが何もできていないので、非常に弱いと思う。

提出したもの

提出物

ソースは以下。コロンはインデントブロックを開始する。ドットは開き括弧と同じだが、対応する閉じ括弧は暗黙で、可能な限り右側で閉じられる。

. \ world0 undocumented -> . let:

  not = . \ x -> . - 1 x
  cadr = . \ x -> . car . cdr x
  cddr = . \ x -> . cdr . cdr x
  caar = . \ x -> . car . car x
  cdar = . \ x -> . cdr . car x
  caddr = . \ x -> . car . cdr . cdr x
  cdddr = . \ x -> . cdr . cdr . cdr x
  caaar = . \ x -> . car . car . car x
  cdaar = . \ x -> . cdr . car . car x
  cadar = . \ x -> . car . cdr . car x
  cddar = . \ x -> . cdr . cdr . car x
  cadddr = . \ x -> . car . cdr . cdr . cdr x
  caaaar = . \ x -> . car . car . car . car x
  cdaaar = . \ x -> . cdr . car . car . car x
  mod = . \ a b -> . - a . * b . / a b
  const = . \ x -> . \ y -> x

  ; create a list that contains @n@ copies of @a@.
  replicate = . \ n a -> .
    if n:
      cons a . replicate (- n 1) a
      0

  ; returns the input list reversed
  reverse = . \ list -> . reverse-go list 0
  reverse-go = . \ list acc -> .
    if (atom list):
      acc
      reverse-go (cdr list) (cons (car list) acc)
  
  ; transform each eleemnt in the list
  map = . \ f list -> . if (atom list):
    list
    cons (f (car list)) (map f . cdr list)

  ; return a list that contains integers m .. n
  enum-from-to = . \ m n -> . if (> m n):
    0
    cons m . enum-from-to (+ m 1) n

  ; length of the list
  length = . \ list -> . if (atom list):
    0
    + 1 . length . cdr list

  maximum-on = . \ default f list -> . if (atom list):
    default
    maximum-on-go f (car list) (f . car list) (cdr list)
  maximum-on-go = . \ f best bestval list -> . if (atom list):
    best
    let:
      val = . f . car list
      if (> val bestval):
        maximum-on-go f (car list) val (cdr list)
        maximum-on-go f best bestval (cdr list)

  minimum-on = . \ default f list -> . if (atom list):
    default
    minimum-on-go f (car list) (f . car list) (cdr list)
  minimum-on-go = . \ f best bestval list -> . if (atom list):
    best
    let:
      val = . f . car list
      if (> bestval val):
        minimum-on-go f (car list) val (cdr list)
        minimum-on-go f best bestval (cdr list)

  ; does the list have more than 1 element?
  list-long-p = . \ list -> .
    if (atom list):
      0
      not . atom . cdr list
  
  pair-list = . \ list -> .
    if (list-long-p list):
      cons (cons (car list) (cadr list)) (pair-list . cddr list)
      if (atom list):
        list
        cons list 0

  ; turn a non-empty list into a perfect binary tree with size
  list-to-bt = . \ list -> .
    list-to-bt-go list 1
  list-to-bt-go = . \ list s -> .
    if (list-long-p list):
      list-to-bt-go (pair-list list) (* s 2)
      cons (car list) s


  ; turn a list of lists into a tree of trees
  list-to-bt2 = . \ list -> . list-to-bt . map list-to-bt list
  list-to-bt3 = . \ list -> . list-to-bt . map list-to-bt2 list

  ; get an item from a binary tree
  index-bt = . \ k bt -> . index-bt-go (car bt) (cdr bt) k
  index-bt-go = . \ tree s k -> .
    if (> s 1):
      do:
        s1 = . / s 2
        if (>= k s1):
          index-bt-go (cdr tree) s1 (- k s1)
          index-bt-go (car tree) s1 k
      tree
  
  ; get an item from a nested binary tree
  index-bt2 = . \ x y bt -> . index-bt x . index-bt y bt

  index-bt3 = . \ x y z bt -> . index-bt z . index-bt2 x y bt

  ; turn a function into a binary tree
  tabulate-bt = . \ sz f -> .
    list-to-bt . map f . enum-from-to 0 (- sz 1)

  ; turn a function into a nested binary tree
  tabulate-bt2 = . \ width height f -> .
    tabulate-bt height . \ y -> .
      tabulate-bt width . \ x -> . f x y

  ; apply @f@ to the positino @k@ of the binary tree @bt@.
  modify-bt = . \ k f bt -> .
    cons (modify-bt-go (car bt) (cdr bt) k f) (cdr bt)
  modify-bt-go = . \ tree s k f -> .
    if (> s 1):
      do:
        s1 = . / s 2
        if (>= k s1):
          cons (car tree) (modify-bt-go (cdr tree) s1 (- k s1) f)
          cons (modify-bt-go (car tree) s1 k f) (cdr tree)
      f tree

  ; modify a nested binary tree
  modify-bt2 = . \ x y f bt -> .
    modify-bt y (\ row -> . modify-bt x f row) bt

  modify-bt3 = . \ x y z f bt -> .
    modify-bt2 x y (\ vs -> . modify-bt z f vs) bt

  ;; skew heap
  ;
  ; heap = 0 | (int, (a, (heap, heap)))

  ; singleton heap
  singleton-heap = . \ prio val -> .
    cons prio . cons val . cons 0 0

  merge-heap = . \ x y -> .
    if (atom x):
      y
      if (atom y):
        x
        if (> (car x) (car y)):
          cons (car y) . cons (cadr y) . cons (merge-heap x (cdddr y)) (caddr y)
          cons (car x) . cons (cadr x) . cons (merge-heap y (cdddr x)) (caddr x)

  insert-heap = . \ p x h -> . merge-heap h . singleton-heap p x

  heap-take-min = . \ h -> .
    cons:
      cons (car h) (cadr h)
      merge-heap (caddr h) (cdddr h)

  ;; queue
  ;
  ; queue = (list, list)

  empty-queue = . cons 0 0

  singleton-queue = . \ x -> . cons (cons x 0) 0

  enqueue = . \ x q -> . cons (car q) (cons x . cdr q)

  dequeue = . \ q -> .
    if (atom (car q)):
      do:
        rev = . reverse (cdr q)
        cons (car rev) (cons (cdr rev) 0)
      cons (caar q) . cons (cdar q) (cdr q)
  
  null-queue = . \ q -> .
    if (atom (car q)):
      atom (cdr q)
      0

  valid-location-p = . \ width height m x y -> . and:
    >= x 0
    > width x
    >= y 0
    > height y
    index-bt2 x y m

  opposite = . \ d -> .
    mod (+ 2 d) 4

  min-ghost-distance-from = . \ width height m start -> . 
    if (== 0 . index-bt2 (car start) (cdr start) m):
      list-to-bt3 . replicate height . replicate width . replicate 4 inf
      min-ghost-distance-from-go width height m:
        modify-bt3 (car start) (cdr start) 0 (const 0) .
          modify-bt3 (car start) (cdr start) 1 (const 0) .
          modify-bt3 (car start) (cdr start) 2 (const 0) .
          modify-bt3 (car start) (cdr start) 3 (const 0) .
          list-to-bt3 . replicate height . replicate width . replicate 4 inf
        enqueue (cons (cons start 1) 0) .
          enqueue (cons (cons start 2) 0) .
          enqueue (cons (cons start 3) 0) .
          singleton-queue (cons (cons start 0) 0)
  min-ghost-distance-from-go = . \ width height m dist q -> . if (null-queue q):
    dist
    do:
      qr = . dequeue q
      x = 0
      y = 0
      z = 0
      next-d = 0
      add = 0
      add-not-opposite = 0
      d1_q1 = 0
      has-way = 0
      
      <- x . caaaar qr
      <- y . cdaaar qr
      <- z . cdaar qr
      <- next-d . + 1 . cdar qr
      <- add . \ z1 dq -> . do:
        x1 = . move-x z1 x
        y1 = . move-y z1 y
        if (valid-location-p width height m x1 y1):
          do:
            <- has-way 1
            if (== inf . index-bt3 x1 y1 z1 dist):
              cons:
                modify-bt3 x1 y1 z1 (const next-d) . car dq
                enqueue (cons (cons (cons x1 y1) z1) next-d) . cdr dq
              dq
          dq
      <- add-not-opposite . \ z1 dq -> . if (== z1 . opposite z):
        dq
        add z1 dq
      <- d1_q1 .
        add-not-opposite 0 .
        add-not-opposite 1 .
        add-not-opposite 2 .
        add-not-opposite 3 .
        cons dist (cdr qr)
      <- d1_q1 . if has-way:
        d1_q1
        add (opposite z) d1_q1
      min-ghost-distance-from-go width height m (car d1_q1) (cdr d1_q1)

  ; calculate distances from a single starting point, @start@.
  min-distance-from = . \ width height m start -> . 
    if (== 0 . index-bt2 (car start) (cdr start) m):
      list-to-bt2 . replicate height . replicate width inf
      min-distance-from-go width height m
        (modify-bt2 (car start) (cdr start) (const 0) .
            list-to-bt2 . replicate height . replicate width inf)
        (singleton-queue (cons start 0))
  min-distance-from-go = . \ width height m dist q -> . if (null-queue q):
    dist
    do:
      qr = . dequeue q
      x = 0
      y = 0
      next-d = 0
      add = 0
      d1_q1 = 0
      
      <- x . caaar qr
      <- y . cdaar qr
      <- next-d . + 1 . cdar qr
      <- add . \ x1 y1 dq -> .
        if (and:
              valid-location-p width height m x1 y1
              == inf . index-bt2 x1 y1 dist):
          cons:
            modify-bt2 x1 y1 (const next-d) . car dq
            enqueue (cons (cons x1 y1) next-d) . cdr dq
          dq
      <- d1_q1 .
        add (- x 1) y .
        add (+ x 1) y .
        add x (- y 1) .
        add x (+ y 1) .
        cons dist (cdr qr)
      min-distance-from-go width height m (car d1_q1) (cdr d1_q1)

  ; minimum distance to the nearest pill
  min-distance-to-pill = . \ width height m start -> . 
    min-distance-to-pill-go width height m
      (modify-bt2 (car start) (cdr start) (const 0) .
          list-to-bt2 . replicate height . replicate width inf)
      (singleton-queue (cons start 0))
  min-distance-to-pill-go = . \ width height m dist q -> . if (null-queue q):
    (- 0 1000) ; no pill is reachable, we have won the game
    do:
      qr = . dequeue q
      x = 0
      y = 0
      next-d = 0
      add = 0
      d1_q1 = 0
      
      <- x . caaar qr
      <- y . cdaar qr
      <- next-d . + 1 . cdar qr
      <- add . \ x1 y1 dq -> .
        if (and:
              valid-location-p width height m x1 y1
              == inf . index-bt2 x1 y1 dist):
          cons:
            modify-bt2 x1 y1 (const next-d) . car dq
            enqueue (cons (cons x1 y1) next-d) . cdr dq
          dq
      <- d1_q1 .
        add (- x 1) y .
        add (+ x 1) y .
        add x (- y 1) .
        add x (+ y 1) .
        cons dist (cdr qr)
      if (== 2 . index-bt2 x y m):
        cdar qr
        min-distance-to-pill-go width height m (car d1_q1) (cdr d1_q1)

  ; minimum distance between each pair of points
  min-distances-among = . \ width height m -> .
    tabulate-bt2 width height . \ x y -> .
      min-distance-from width height m (cons x y)

  ; infinite distance
  inf = 100000

  make-sim-state = . \ width height m loc score -> .
    cons:
      cons:
        cons width height
        cons m loc
      score
  state-width = . \ state -> . caaar state
  state-height = . \ state -> . cdaar state
  state-map = . \ state -> . cadar state
  state-location = . \ state -> . cddar state
  state-score = . \ state -> . cdr state

  move-x = . \ dir x -> .
    if (== dir 1):
      (+ x 1)
      if (== dir 3):
        (- x 1)
        x

  move-y = . \ dir y -> .
    if (== dir 0):
      (- y 1)
      if (== dir 2):
        (+ y 1)
        y

  movable = . \ dir state -> .
    valid-location-p:
      state-width state
      state-height state
      state-map state
      move-x dir . car . state-location state
      move-y dir . cdr . state-location state

  simulate-move = . \ dir state -> . do:
    new-x = . move-x dir . car . state-location state
    new-y = . move-y dir . cdr . state-location state
    make-sim-state:
      state-width state
      state-height state
      modify-bt2 new-x new-y (const 1) . state-map state
      cons new-x new-y
      +:
        state-score state
        if (== 2 . index-bt2 new-x new-y . state-map state):
          10
          0

  test-map = 0
  current-map = 0
  map-width = 0
  map-height = 0
  
  map-distances = 0
  directions = . list 0 1 2 3

  distance-to = . \ from to -> .
    index-bt2 (car to) (cdr to) .
      index-bt2 (car from) (cdr from) map-distances

  think = . \ state world -> . let:
    man-state = . cadr world
    ghost-states = . caddr world
    enemy-distance-score = . \ state -> .
      minimum-on 0 (\ x -> x) .
        map (\ ghost -> . distance-to (state-location state) . cadr ghost)
        ghost-states
    score-state = . \ state -> . +:
      state-score state
      -:
        * 2 . enemy-distance-score state
        min-distance-to-pill:
          state-width state
          state-height state
          state-map state
          state-location state

    ;score-move = . \ dir -> . if (movable dir sim-state):
    ;  score-state . simulate-move dir sim-state
    ;  0
    score-move = . \ dir -> . do:
      score = . if (movable dir sim-state):
        score-state . simulate-move dir sim-state
        0
      trace (list 1 dir score) score
    man-location = 0
    sim-state = 0
    next-direction = 0

    do:
      trace (cons 5 ghost-states) .
        <- man-location . cadr man-state
      <- sim-state . make-sim-state map-width map-height current-map man-location 0
      <- next-direction . maximum-on 100 score-move directions
      <- current-map . state-map . simulate-move next-direction sim-state
      trace (cons 0 next-direction) .
        cons state next-direction

  do:
    <- current-map . list-to-bt2 . car world0
    <- map-width . length . car . car world0
    <- map-height . length . car world0
    <- map-distances .
      min-distances-among map-width map-height current-map

    <- test-map . list-to-bt2 . list:
      list 1 0 1
      list 1 1 1
    ; step = . \ state world -> . cons 0 1
    ; cons 0 step
    ;index-bt (list-to-bt . cons 100 . cons 1001 . cons 1002 0) 1
    ; heap-take-min . cdr . heap-take-min . insert-heap 1 100 . insert-heap 3 300 . insert-heap 2 200 . singleton-heap 9 900
    ; update-bt 3 8 (list-to-bt . replicate 10 4)
    ; dequeue . cdr . dequeue . enqueue 1 . enqueue 2 . enqueue 3 . singleton-queue 4
    ; modify-bt2 0 1 (\ x -> . + x 1) .
    ; modify-bt2 0 1 (\ x -> . + x 1) . cons (cons (cons (cons 0 1) 2) (cons (cons 2 3) 2)) 2
    ; min-distances-among 3 2 test-map
    ;trace (min-ghost-distance-from 3 2 test-map (cons 0 0)) 0

    cons 12345 think
    ;think 12345 world0

*1:パックマンらしい。未プレイ

*2:数時間で作るので必然的に糞言語になる

*3:ダイクストラ法を書くつもりだったが結局使わなかった

ノイズ耐性のある二分探索

観測にノイズが乗っても対処できる二分探索が意外と簡単に書けることが分かったのでメモ。

問題

(n-1)個の整数からなる列がある。そのうち左からi個は(-1)であり、それ以外は1である。

n=6, i=3の例
-1 -1 -1 1 1

クエリを繰り返すことでiを求めたい。各クエリは整数kであり、答として左からk番目の整数の値が得られる。ただし、この答にはノイズが加算される。ノイズは標準偏差σ(既知とする)の正規分布に従う。

解法

iがどの値を取るかの確率分布を持っておいて、クエリの答が得られるたびにベイズの定理に従って更新する。最初はn通りの一様分布。クエリは、得られる情報量の期待値を最大化するように選ぶ。これには、(-1)と1の境界よりも右にあるか左にあるかが半々に近い位置を選べば良い。

iの確率分布が十分に偏ったら終了。

実験

以下では、iが特定の値を取る確率が95%を越えた時点で探索を終了するようにしている。

ノイズが小さい場合は普通の二分探索のように動作する。

# n=8, σ=0.3, i=3
% ./noisy-bsearch 8 0.3 --answer 3 
? 3 -> 1.031
? 1 -> -0.579
? 2 -> -0.859
Found: 3 posterior=0.9999999944164625

ノイズが大きくなると、解を確定するのに同じクエリを何度も発行する必要が出てくる。

# n=8, σ=1, i=3
% ./noisy-bsearch 8 1 --answer 3
? 3 -> 1.214
? 1 -> -0.943
? 2 -> 0.184
? 2 -> 0.901
? 1 -> -0.704
? 2 -> -2.098
? 3 -> 1.097
? 2 -> -1.715
Found: 3 posterior=0.9579616340601826

たまに間違える。

% ./noisy-bsearch 8 1 --answer 3
? 3 -> 1.813
? 1 -> 1.609
? 0 -> -1.378
? 1 -> 0.376
? 0 -> -1.202
Found: 1 posterior=0.9565230889174586

nを大きくすると、適当に当たりを付けながら試行錯誤する様子が見える。

# n=100, σ=2, i=77
% ./noisy-bsearch 100 2 --answer 77
? 49 -> -1.182
? 60 -> -1.373
? 70 -> -0.116
? 71 -> 0.242
? 69 -> -0.160
? 70 -> -4.324
? 83 -> -0.343
? 84 -> -0.377
? 85 -> 2.448
? 79 -> 4.538
? 73 -> -0.290
? 74 -> -0.890
? 75 -> -1.684
? 77 -> 2.514
? 75 -> -1.450
? 76 -> -1.166
? 76 -> -4.089
? 77 -> 0.621
? 77 -> -2.006
? 78 -> 2.714
? 77 -> -0.352
? 77 -> 1.402
? 77 -> -0.629
? 77 -> -2.606
? 77 -> -0.347
? 78 -> 1.579
? 77 -> 7.286
? 76 -> 2.773
? 76 -> 0.334
? 76 -> -1.448
? 76 -> 0.915
? 76 -> 1.719
? 76 -> -2.039
? 76 -> -3.672
? 76 -> -0.300
? 77 -> 1.227
? 76 -> -3.140
Found: 77 posterior=0.9527731713033516

まとめ

推定すべき確率分布が離散的だと、頑張って近似しなくても教科書通りのベイズの定理が使えて楽しいことが分かった。

コード

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE TemplateHaskell #-}

import Control.Applicative
import Control.Monad
import Control.Lens
import Data.Function
import Data.List
import Data.Maybe
import qualified Data.Vector.Unboxed as U
import Statistics.Distribution
import Statistics.Distribution.Normal
import System.Environment
import System.Console.GetOpt
import System.IO
import qualified System.Random.MWC as MWC
import Text.Printf
import Text.Read

data Interaction = Human | Simulated !Int
data Mode = Help | Search

data Conf = Conf
  { _observationFile :: !(Maybe FilePath)
  , _interaction :: !Interaction
  , _verbose :: !Bool
  , _mode :: !Mode
  }

makeLenses ''Conf

options :: [OptDescr (Conf -> Conf)]
options =
  [ Option "f" ["observation-file"]
      (ReqArg (\f -> observationFile .~ Just f) "FILE")
      "Apply observations from FILE before starting search"
  , Option "a" ["answer"]
      (ReqArg (\a -> interaction .~ Simulated (read' a)) "N")
      "Automatically respond to queries"
  , Option "v" ["verbose"]
      (NoArg (verbose .~ True))
      "Display posterior distribution at each step"
  , Option "h" ["help"] (NoArg (mode .~ Help)) "Show this help"
  ]

defaultConf :: Conf
defaultConf = Conf
  { _observationFile = Nothing
  , _interaction = Human
  , _verbose = False
  , _mode = Search
  }

read' :: (Read a) => String -> a
read' s = fromMaybe (error $ "cannot parse " ++ s) $ readMaybe s

main :: IO ()
main = do
  args <- getArgs
  either fail id $ do
    let !(fns, nonOptions, errs) = getOpt Permute options args
    case errs of
      [] -> return ()
      _ -> Left $ unlines errs
    let conf = foldr ($) defaultConf fns
    case conf^.mode of
      Help -> return $
        putStr $ usageInfo "noisy-bsearch N SIGMA" options
      Search -> do
        (n, sigma) <- case nonOptions of
          [read' -> n, read' -> sigma] -> return (n, sigma)
          _ -> Left "expecting 2 arguments"
        return $ doSearch conf n sigma 

doSearch :: Conf -> Int -> Double -> IO ()
doSearch conf n sigma = do
  obss <- case conf^.observationFile of
    Nothing -> return []
    Just path -> map parse . lines <$> readFile path
  getAns <- getAnswer sigma $ conf^.interaction
  search conf sigma getAns
    $ foldl' (update sigma) (uniform n) obss
  where
    parse (words -> [a, b]) = (read a, read b)
    parse str = error $ "parse errro: "  ++ str

getAnswer :: Double -> Interaction -> IO (Int -> IO Double)
getAnswer _ Human = return askInteractive
getAnswer sigma (Simulated ansIdx) = do
  gen <- MWC.createSystemRandom
  return $ \k -> do
    r <- genContVar
      (normalDistr (if ansIdx <= k then 1 else -1) sigma) gen
    _ <- printf "? %d -> %.3f\n" k r
    return r

askInteractive :: Int -> IO Double
askInteractive k = do
  putStr $ show k ++ "? "
  hFlush stdout
  ans <- getLine
  case readMaybe ans of
    Just r -> return r
    Nothing -> do
      putStrLn "Try again"
      askInteractive k

-- | 実際の探索ループ
search :: Conf -> Double -> (Int -> IO Double) -> Dist -> IO ()
search conf sigma ask = loop
  where
    loop dist
      | Just k <- U.findIndex (>0.95) dist =
        printf "Found: %d posterior=%f\n" k (dist U.! k)
      | otherwise = do
        when (conf ^. verbose) $ print dist
        let k = medianIndex dist
        obs <- ask k
        loop $ update sigma dist (k, obs)

-- | いくつかの選択肢から一つを選ぶ確率分布
type Dist = U.Vector Double

-- | 一様分布
uniform :: Int -> Dist
uniform n = U.replicate n (1 / fromIntegral n)

-- | 観測データと事前分布から事後分布を得る
update :: Double -> Dist -> (Int, Double) -> Dist
update sigma dist (k, obs) = normalize $ U.imap upd dist
  where
    upd i x
      | i <= k = x * likR
      | otherwise = x * likL
    -- kが正解の右側にある場合の尤度
    likR = exp (- ((obs - 1) / sigma) ^ (2::Int) / 2) / (sigma * pi * sqrt 2)
    -- kが正解の左側にある場合の尤度
    likL = exp (- ((obs + 1) / sigma) ^ (2::Int) / 2) / (sigma * pi * sqrt 2)

-- | P(answer <= k) が0.5に最も近いkを返す
medianIndex :: Dist -> Int
medianIndex = U.minIndexBy (compare `on` (abs . subtract 0.5)) . U.scanl1 (+)

-- | 合計が1になるように定数倍する
normalize :: U.Vector Double -> Dist
normalize xs = U.map (/U.sum xs) xs