ABC 141 E - Who Says a Pun? : Z-Algorithm を Javaで実装する

E - Who Says a Pun?の復習。

解説の中にあるZ-Algorithmを勉強したので記録。

こちらのリンクが大変ためになった。 www.geeksforgeeks.org

Z-Algorithmとは

Z-Algorithmとは、『Z-Array』を線形時間で構築するアルゴリズムであり、ある文字列に対して、あるパターンが出現するかどうかを判定する際などに用いられる。

Z-Arrayとは「文字列str[0..n-1]が与えられた時、Z[i]は文字列str[0..n-1]とstr[i..n-1]の最長共通接頭辞数であるような配列」である。

Example:
Index            0   1   2   3   4   5   6   7   8   9  10  11 
Text             a   a   b   c   a   a   b   x   a   a   a   z
Z Array          X   1   0   0   3   1   0   0   2   2   1   0 

例えばここでパターンを "hoge" とおいて、テキストを "fugahogepiyo" としたとき、"hoge$fugahogepiyo" *1に対してZ-Algorithmを適用しZ-arrayの中にパターンの文字列長の値が入っているかを確かめることで、テキスト中にパターン "hoge" が登場するかどうかを判定することができる。

O(n)で解くために、計算途中での共通接頭辞をメモしているのがキモ。 共通接頭辞の始まるindexをL、最後の文字のindexをRとしたとき、str[0..R-L]とstr[L..R]は同じ文字列であることを利用して、計算量を削減している。 つまり、i<Rである時、ナイーブに共通接頭辞を見るのではなくて、すでに計算されたZ[i - L]を利用する。

解答

import java.io.*;
import java.util.*;

class Main {

    public static void main(String[] args) throws Exception {
        final Scanner sc = new Scanner(System.in);
        int N;
        N = sc.nextInt();
        String S;
        S = sc.next();
        solve(N, S);
    }

    static void solve(int N, String S) {
        int ans = 0;
        for (int i = 0; i < N; i++) {
            int[] z = z_algorithm(S.substring(i));
            for (int j = i; j < N; j++) {
                ans = Math.max(ans, Math.min(j - i, z[j - i]));
            }
        }
        System.out.println(ans);

    }

    static int[] search(String pattern, String text) {
        return z_algorithm(pattern + "$" + text);
    }

    static int[] z_algorithm(String S) {
        int l = 0;
        int r = 0;
        int[] Z = new int[S.length()];

        for (int i = 1; i < S.length(); i++) {
            if (r < i) {
                l = r = i;
                while (r < S.length() && S.charAt(r) == S.charAt(r - l)) {
                    r++;
                }
                Z[i] = r - i;
                r--;
            } else {
                int k = i - l;
                if (i + Z[k] - 1 < r) {
                    Z[i] = Z[k];
                } else {
                    l = i;
                    while (r < S.length() && S.charAt(r) == S.charAt(r - l)) {
                        r++;
                    }
                    Z[i] = r - i;
                    r--;
                }
            }
        }
        Z[0] = S.length();
        return Z;
    }

*1:$はパターン内にもテキスト内にも出現しないことが保証されているとする。