Module f

Expand description

§ABC 386 F - Operate K

refs: https://atcoder.jp/contests/abc386/tasks/abc386_f

レーベンシュタイン距離を利用する. レーベンシュタイン距離は, 文字の挿入, 削除, 置換の操作を用いて文字列を変換する最小の操作数を表す. 文字列 $S$ ($|S| = N$)と文字列 $T$ ($|T| = M$)のレーベンシュタイン距離は次の DP で求められる.

// dp[i][j]: Sのi文字目までの文字列とTのj文字目までのレーベンシュタイン距離

for i in 0..N:
  dp[i][0] = i  // Sのi文字目までの文字列を空文字列に変換するにはi回削除することが自明

for j in 0..M:
  dp[0][j] = j  // 空文字列をTのj文字目までの文字列に変換するにはj回挿入することが自明

for i in 1..N:
  for j in 1..M:
    dp[i][j] = min(
        dp[i-1][j] + 1,  // 文字削除
        dp[i][j-1] + 1,  // 文字挿入
        dp[i-1][j-1] + (S[i-1] != T[j-1] ? 1 : 0)  // 文字置換
    )

return dp[N][M]

この DP を行う場合, オーダーは $O(NM)$となるが, 今回の問題では文字列の長さが最大で $5 \times 10^5$ となるため, このままでは間に合わない. $i$ 文字目までの文字列と $j$ 文字目までの文字列のレーベンシュタイン距離を求める必要性はあるが, $i$ と $j$ の差が $k$ より大きくなる部分は, 文字の削除 / 挿入の回数が $k$ を超えるため, 確定的に値が $k$ を超過し, 今回の問題に置いては計算の必要性はない.

このため, 各 $i$ について, $j$ の範囲を $i - k$ から $i + k$ までに限定して DP を行うことで, 計算量を削減することができる. 探索する範囲を絞る都合上, DP の添字が参照する $i$ の階層によってズレる点に注意.

§元の DP
012345
i = 0j = 0j = 1j = 2j = 3j = 4j = 5
i = 1j = 0j = 1j = 2j = 3j = 4j = 5
i = 2j = 0j = 1j = 2j = 3j = 4j = 5
§範囲絞り DP (ex. k = 2 の場合)
012345
i = 0j = -2j = -1j = 0j = 1j = 2j = 3
i = 1j = -1j = 0j = 1j = 2j = 3j = 4
i = 2j = 0j = 1j = 2j = 3j = 4j = 5
// dp[i][d]: S の i 文字目までの文字列と T の(i + d - k)文字目までのレーベンシュタイン距離

for i in 0..min(N, k):
dp[i][k - i] = i //  T の 0 文字目に当たる部分が i の値に応じてズレる

for d in 0..min(M, k):
dp[0][d + k] = d // T の 0 ~ k 文字目に対する初期化のため, i = 0 の場合 k ~ 2k を初期化する

for i in 1..N:
    for d in 0..(2 * k):
        j = i + d - k
        if j < 0 or j > M:
            continue // j が範囲外の場合はスキップ

        dp[i][d] = min(
            dp[i-1][d + 1] + 1,  // 一つ前のiの階層を参照するため, +1 dがズレる
            dp[i][d-1] + 1,      // 同一階層のため元の処理と同一
            dp[i-1][d] + (S[i-1] != T[j-1] ? 1 : 0)  // 一つ前のiの階層を参照するため, +1 dがズレる
        )

// T の M 文字目となる d (M = N + d - k => d = M + k - N)
return dp[N][M + k - N]
use proconio::{input, marker::Chars};

const INF: usize = 1 << 60;

fn main() {
    input! {
        k: usize,
        s: Chars,
        t: Chars,
    }

    if s.len().abs_diff(t.len()) > k {
        println!("No");
        return;
    }

    let distance = levenshtein_distance(k, &s, &t);

    if distance <= k {
        println!("Yes");
    } else {
        println!("No");
    }
}

fn levenshtein_distance(k: usize, s: &[char], t: &[char]) -> usize {
    let n = s.len();
    let m = t.len();
    let mut dp = vec![vec![INF; 2 * k + 1]; n + 1];

    for i in 0..=k.min(n) {
        dp[i][k - i] = i;
    }
    for d in 0..=k.min(m) {
        dp[0][d + k] = d;
    }

    for i in 1..=n {
        for d in 0..=(2 * k) {
            let j = i as isize + d as isize - k as isize;
            if j < 0 || j > m as isize {
                continue;
            }

            let j = j as usize;
            if d < 2 * k {
                dp[i][d] = dp[i][d].min(dp[i - 1][d + 1] + 1); // 削除
            }

            if d > 0 {
                dp[i][d] = dp[i][d].min(dp[i][d - 1] + 1); // 挿入
            }

            if j > 0 {
                let cost = if s[i - 1] == t[j - 1] { 0 } else { 1 };
                dp[i][d] = dp[i][d].min(dp[i - 1][d] + cost); // 置換
            }
        }
    }

    dp[n][m + k - n]
}