TCBの過去問紹介【TechFUL上級問題にチャレンジ！】

あいさつ

こんにちはこんばんは。 TechFULでアルバイトをしている あるまかん(@Arumakan_ei1727) です。
今回は私がTechFULで作問した問題を1問紹介したいと思います。

問題設定はとてもシンプルなので、ぜひ挑戦してくださるとありがたいです！

問題紹介

紹介する問題は 「最良のキャッチコピー」 という問題です。
2020年10月のTCBで出題されました。

ちなみにTechFULでは毎月 TCB (TechFUL Coding Battle) というプログラミングコンテストを開催しています。 ぜひ参加してください！！

TCBについてはこちら

問題概要

問題概要は次のとおりです。

• 長さ $N$ の文字列 $S_0$ と整数 $Q$ が与えられる。

• 以下の $Q$ 回の操作を行って、新たに文字列 $S_1, \ldots , S_Q$ を作る。

  • $i \ (1 \le i \le Q)$ 回目の操作では2つの整数 $X_i, Y_i$ が与えられる。

  • $S_{i-1}$ の $X_i$ 文字目と $Y_i$ 文字目を入れ替えてできる文字列を $S_i$ とする
    ($S_{i-1}$ に変更は加えない) 。

• $Q+1$ 個の文字列 $S_0, S_1, \ldots , S_Q$ の中で辞書順最小の文字列を求めてください。

※実際の問題文では、「$Q+1$ 個の文字列 $S_0, S_1, \ldots , S_Q$ の中で辞書順最小の文字列」が「最良のキャッチコピー」と定義されています。これが問題タイトルの由来です。

制約

入力される値は以下の条件を満たすことが保証されています。

• $2 \le N \le 10^5$
• $1 \le Q \le 10^5$
• $1 \le X_i \lt Y_i \le N$
• $S_0$ は英小文字のみからなる長さ $N$ の文字列

実行時間制限は 5.0秒 です。

例

例として下記のケースを挙げます。

• $S_0 = $ cabx, $N = 4$
• $Q = 2$
• $(X_1, Y_1) = (1, 3)$
• $(X_2, Y_2) = (2, 3)$

1回目の操作で $S_0$ の 1文字目と3文字目がスワップされて $S_1 = $ bacx が生成されます。
2回目の操作で $S_1$ の 3文字目と4文字目がスワップされて $S_2 = $ baxc が生成されます。

$S_0, S_1, S_2$ のうちで辞書順最小の文字列は $S_1$ の bacx であり、これが答えです。

以下ネタバレ注意

ここから下は解法のネタバレが含まるので注意です。
解き方が分かったら、ぜひ実際にプログラムを書いてコードを提出してみてください。

愚直解法

以下、ある文字列 $A$ の左から $j$ 文字目を $A[j]$ と表記することにします。

まず素直な解法として、問題文のとおりにシミュレーションしていく方法が考えられます。
現時点で辞書順最小の文字列を保持しておく変数 $T$ を設けて、$S_i$ 中の2文字を入れ替えるたびに $T$ と辞書順比較して $T$ を更新していく戦略です。

例えば Python3 なら以下のようなコードとなるでしょう。

N, Q = map(int, input().split())
S = list(input())  # Pythonでは文字列は変更できないので文字のリストにする

T = S.copy()  # 現時点で辞書順最小のS

for _ in range(Q):
    x, y = map(int, input().split())
    x -= 1
    y -= 1
    S[x], S[y] = S[y], S[x]
    if S < T:
        T = S.copy()

print("".join(T))

この解法は TLE(実行時間制限オーバー) になります。

$Q$ 個のスワップ操作を処理するのにちょうど $Q$ 回のループが必要です。
$Q$ 回の各ループの中では $S$ と $T$ の辞書順比較を行っています。辞書順比較は最悪で $N$ 回の比較を行うことになるので、プログラム全体の計算量は $O(NQ)$ となります。

さらに言えば、$T$ に $S$ をコピーする操作でも $\Theta(N)$ かかってしまいます。

制約の項で述べたとおり、$N$ も $Q$ も最大で $10^5$ の値をとります。
$10^5 \times 10^5$ 回の処理は実行時間制限の5.0秒内では間に合わず、TLE になります。

考察

$X_i, Y_i$ の入力もあり、$Q$ 回のループは避けられなさそうです。
どうにかして、もうひとつのボトルネックである辞書順比較を高速に行いたくなります。

また、1回の操作につき $S$ が高々2文字しか変化しないという問題の特性も鍵になりそうです。

辞書順比較ではどこを比較しているか？

例えば、2つの文字列 iiiabc と iiixyz を辞書順比較してみてください。 おそらく a と x を比較したかと思います。
i 同士の比較や、 b と y との比較などは飛ばしたはずです。

辞書順比較の別の定義

先ほどの考察から、文字列 $S, T$ の辞書順比較は次のように言い換えられます。

• $S[j] \ne T[j]$ であるような文字の組 $(S[j], T[j])$ のうち、 $j$ が最小であるような組の文字を比較する。
  • そのような組がない場合は、比較結果は $S = T$ 。

この比較方法は、 $S[j]$ と $T[j]$ の1文字のみの比較なので高速です。

差分情報の構築

先ほどの1文字比較を実現するために、まず $(S[j], T[j])$ の組を全て求めなければなりません (全て求めたあとに、$j$ が最小であるような組を採用して文字比較を行います)。
これらの組の集合は $S, T$ の相異なる文字の組の集合なので、 $S$ と $T$ の間の差分情報と言えます。

さて、差分情報を求めるために毎回 $j$ を $1$ から $N$ までループしてしまうと、 全体で $\Theta(NQ)$ となって間に合いません。

ここで、スワップ操作の前後で差分情報は高々2文字しか変化しないことに注目します。
$j = X_i, Y_i$ の位置のみチェックすれば十分ですね。
$S[X_i] \ne T[X_i]$ かどうか調べて、必要に応じて差分情報の集合に追加・削除をすればいいです ($Y_i$ についても同様) 。

差分情報の表現

差分情報は $S[j] \ne T[j]$ であるような整数 $j$ の集合としてよいです。

集合への追加・削除と、集合内の最小要素の取得は、平衡二分探索木を用いた順序付き集合を用いれば $O(\log N)$ で処理できます (集合に追加される要素数は最大で $N$ 個)。
C++ の std::set や Java の java.util.SortedSet (TreeSet) は順序付き集合です。

Python の標準ライブラリには順序付き集合の実装がありませんが、今回の解法はハッシュセットの dict と 二分ヒープの heapq を併用すれば実現でき、$O(\log N)$ です。

辞書順比較が $O(\log Q)$ でできました！🎉🎉🎉🎉🎉🎉🎉

$T$ の更新の計算量改善

前述の愚直解の Python3 コードにおいて、S < T の比較部分は計算量を改善できました。
T = S.copy() の部分はどうすればよいでしょうか？ $N$ 文字全てコピーしていては全体で $\Theta(NQ)$ となり間に合いません。

コピーが $Q$ 回発生する例としては、$S_0 = $ zxy...cba zxy...cba ... で、スワップするたびに必ず辞書順で小さくなるような操作列が与えられた場合が考えられます。 $Q$ 回の各スワップで $T$ の更新が発生するため、$\Theta(NQ)$ となりますね。

# 愚直解プログラムの抜粋
    if S < T:
        T = S.copy()

$S$ の $N$ 文字全てをコピーする必要はありません。
$T$ に適用されていないスワップ操作を、順に $T$ に適用していけば $T = S$ にできます。

$T$ に適用するスワップ操作の総数は、$Q$ 回のループ全体をとおして最大で $Q$ 回なので間に合います。

※補足しておくと、T = S と代入した場合は T は S の参照を持つことになるので、S が変更されると T も変更されたかのようなことになってしまいます。

答え

文字列比較が $O(\log N)$ 、文字列のコピーがならし $O(\log N)$ ででき、全体で $O(Q \log N)$ で解けました。

以下に C++ と Python3 による解答例を載せておきます。

なお実際の問題で出力すべき値は、文字列 $T$ ではなく $S_0, \ldots , S_Q$ の中で辞書順最小の文字列が $S_k$ であるような添字 $k$ であることに注意してください。 (出力量が多すぎてTLEになることを防ぐためです。今回の問題紹介では非本質だったのでこれまで省略してきました。)

C++ 解答例 (クリック・タップして開きます)

差分情報は添字 $j$ のみを保持する std::set で表現しても良かったのですが、このプログラムでは $j$ と $S[j]$ を保持するようにしています。 $j$ だけでなく $S[j]$ も保持することで、$T$ に差分を適用する処理を簡単にしています。

#include <iostream>
#include <map>
#include <string>
using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, Q;
    string S;
    cin >> N >> Q >> S;

    string T = S; // 現時点で辞書順最小のS
    int ans_index = 0; // S_0, ... S_Q 内で S_k が辞書順最小となるような添字 k
    map<int, char> diffs; // SとTの差分。S[j] != T[j] のとき diffs[j] = S[j] 

    for (int i = 0; i < Q; ++i) {
        int x, y;
        cin >> x >> y;
        --x, --y;
        swap(S[x], S[y]);

        if (S[x] == T[x]) diffs.erase(x);
        else diffs[x] = S[x];

        if (S[y] == T[y]) diffs.erase(y);
        else diffs[y] = S[y];

        if (diffs.empty()) continue; // 差分がない => S == T なので T を更新する必要はない

        const int min_j = diffs.begin()->first;
        if (S[min_j] < T[min_j]) {  // S < T なら T に差分を適用して T == S の状態にする
            ans_index = i + 1;
            for (const auto& [j, chr] : diffs) {
                T[j] = chr;
            }
            diffs.clear();
        }
    }

    cout << ans_index << endl;

    return 0;
}

Python 解答例 (クリック・タップして開きます)

def main():
    import sys
    import heapq
    input = sys.stdin.readline

    N, Q = map(int, input().split())
    S = list(input().strip())

    T = S.copy()  # 現時点で辞書順最小のS
    ans_index = 0  # S_0, ... S_Q 内で S_k が辞書順最小となるような添字 k
    diffs = dict()  # SとTの差分。S[j] != T[j] のとき diffs[j] = S[j]
    minheap_j = list()  # S[j] != T[j] であるような j のうちで最小の j を求めるためのMinヒープ

    for i in range(Q):
        x, y = map(int, input().split())
        x -= 1
        y -= 1
        S[x], S[y] = S[y], S[x]

        if x not in diffs:
            heapq.heappush(minheap_j, x)
        if y not in diffs:
            heapq.heappush(minheap_j, y)
        diffs[x] = S[x]
        diffs[y] = S[y]

        # S[j] == T[j] であるような j が含まれている可能性があるので削除
        while len(minheap_j):
            min_j = minheap_j[0]
            if S[min_j] != T[min_j]:
                break
            heapq.heappop(minheap_j)
            diffs.pop(min_j)

        # 差分がない => S == T なので T を更新する必要はない
        if len(minheap_j) == 0:
            continue
        # S < T なら T に差分を適用して T == S の状態にする
        min_j = minheap_j[0]
        if S[min_j] < T[min_j]:
            ans_index = i + 1
            for j, ch in diffs.items():
                T[j] = ch
            diffs.clear()
            minheap_j.clear()

    print(ans_index)


if __name__ == "__main__":
    main()

問題の本質と応用

私はこの問題の本質は以下の2点であると考えます。

• 操作の前後で、変化する箇所が限定されている
• 辞書順比較における無駄の削減

他の問題を解く際も、以上のような点に注意すると解法が見えやすくなるかもしれません。
特に1つ目の「変化する箇所が限定されている」かどうかは多くの問題に当てはめることができるかと思います。

また、ある2つの文字列 $A, B$ の最長共通接頭辞を $\mathrm{LCP}(A, B)$ とすると、$A$ と $B$ の辞書順比較は
$A[\mathrm{LCP}(A, B) + 1]$ と $B[\mathrm{LCP}(A, B) + 1]$
の1文字比較であるという見方もできますね。
知っておくとどこかで役に立つかもしれません。

終わりに

ところで、あるまかんは今は作問というよりは開発アルバイトをメインにしています。
開発も楽しいです！
もしかしたら、またいつか作問してコンテストに出題するかもしれません (???)

今回紹介したのはTCBの過去問です。 TechFULでは、TCBの過去問などはプラクティス問題として公開されています。
プラクティス問題は、普通の問題とは違ってテストケースが公開されており、Twitterで解法ツイートをしたりしても大丈夫です！

TechFULにはもう一つ、チャレンジ問題というジャンルがあります。
こちらはスキル測定用なので解法ツイートはできませんが、幅広いジャンルの問題がたくさんあります。
解いた分だけ確かな実力として記録が残ります。 ぜひ挑戦してみてください！

最後まで読んでくださりありがとうございました！