【ABC452-C】前処理でTLEを回避するアプローチ

記事の目的

クエリをまたいで変わらない計算を事前に済ませることを理解する。

最初の提出ではTLEではねられました（テストケースは2つとも通ったのに）。ここ数週間はAtCoderに精力的に参加していますが（編入対策の一環として）、定期的に計算量を考えることの必要性に気付かされます。その点についてもまとめていきたいです。

問題の概要

$N$本の肋骨と1本の脊椎を持つオブジェがあります。$M$個の文字列の中から、各骨に1つずつ割り当てます。脊椎候補$S_j$ごとに条件を満たす割り当てが存在するかを答えます。

肋骨$i$の条件は以下の通りです。

• 割り当てる文字列の長さが$A_i$
• その$B_i$文字目が脊椎の$i$文字目に一致する

最初の実装（TLE）

設計として各クエリ$S_j$に対して、全肋骨×全文字列を毎回探索します。

for (int j = 0; j < M; j++) {
    if (strlen(S[j]) != N) {
        printf("No\n");
        continue;
    }
    
    int count = 0;
    for (int i = 0; i < N; i++) {
        for (int k = 0; k < M; k++) {
            if (strlen(S[k]) == A[i] && S[k][B[i] - 1] == S[j][i]) {
                count++;
                break;
            }
        }
    }
    if (count == N) printf("Yes\n");
    else printf("No\n");
}

計算量を考えてみましょう。

O(M \times N \times M) = O(M^2N)

$M=200{,}000$、$N=10$のため計算量は約$4 \times 10^{11}$回となり、タイムアウトします。

最適化（AC）

「肋骨$i$に文字$c$が来たとき、条件を満たす文字列が存在するか」という答えは、どのクエリでも変わりません。クエリ書えりの前に1回だけ計算しておけます。

2次元配列can[10][26]を宣言し、can[i][c] = 肋骨$i$に文字$c$（a ~ z）が来たとき使える文字列が存在するかを保存します。

// 前計算
for (int i = 0; i < N; i++) {
    for (int k = 0; k < M; k++) {
        if (strlen(S[k]) == A[i]) {
            can[i][S[k][B[i] - 1] - 'a'] = 1;
        }
    }
}

// クエリ処理
for (int j = 0; j < M; j++) {
    if (strlen(S[j]) != N) {
        printf("No\n");
        continue;
    }
    int count = 0;
    for (int i = 0; i < N; i++) {
        if (can[i][S[j][i] - 'a']) count++;
    }
    if (count == N) printf("Yes\n");
    else printf("No\n");
}

最適化後の計算量を考えてみましょう。

前計算で$O(N \times M)$、クエリ処理で$O(M \times N)$なので合計$O(N \times M)$になります。$N=10$、$M=200{,}000$のため計算量は約200万回となりACします。

おわりに

高校受験では「ひたすらたくさん問題を解く！」という方針でやっていたので編入勉強では「1問1問丁寧に理解する」という方針で精進していきたいです。