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

背景

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

例として、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]
}

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

経路復元が必要な場合

ここからが本題。

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

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

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

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

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

レーベンシュタイン距離の問題では、二つの文字列の長さを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()
}

なぜこれを書いたか

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