Skip to content

1397. Find All Good Strings

#188

Problem link

https://leetcode.com/problems/find-all-good-strings/submissions/

Problem Summary

사전순으로 s1 와 s2 사이의 n 길이 문자열 중에 evil 이 부분 문자열로 안 들어간 개수를 구하는 문제.

Solution

꽤나 어렵다.

https://leetcode.com/problems/find-all-good-strings/discuss/554806/O(N-*-len(evil))-Solution-with-memorized-DP-and-KMP-algorithm
참고함.

먼저 DP라는 건 modular 연산이나 감으로 알 수 있다. s1, s2 사이의 문자열의 개수를 DP로 구현할 수 있다.
문제는 evil 문자열을 걸러내는 건데... KMP 알고리즘을 사용하면 된다.

KMP 는 로이형의 유튜브를 참고
https://www.youtube.com/watch?v=GTJr8OvyEVQ

본격적인 코드 설명에 들어가보면...
일단 s1, s2의 경계를 알기 위해 is_prefix_s1, is_prefix_s2 변수를 사용하고 현재까지 나온 evil 인덱스를 prefix_evil으로 설정하였다.

first, last의 값은 탐색할 문자가 되고 is_prefix_s1, is_prefix_s2 경계에 위치한 경우 경계 바깥은 탐색할 필요가 없으므로 빼줘야 한다.
이제 evil 부분 문자열 체크가 있는데 KMP 알고리즘을 사용한다. KMP 알고리즘을 이해하고 있으면 쉬운데 간단하게 설명하자면 탐색할 다음 인덱스를 구하는 것이다.
즉, evil 문자열이 aab 라고 할 때 aaab 문자열의 경우 aaa 탐색 후 aaab가 evil 문자열인지를 체크하려면 다시 처음부터 봐야 한다. 여기서 KMP 알고리즘의 next 배열을 사용하면 처음부터가 아닌 aa는 이미 매치된 것이고 뒤에 b 하나만 나오면 evil이랑 매치가 된다고 보는 방식이다. 이렇게 하면 하나씩 탐색하면서 부분문자열인지 판별 가능하다.

Source Code

from functools import lru_cache


class Solution:
    def findGoodStrings(self, n: int, s1: str, s2: str, evil: str) -> int:

        mod = 10 ** 9 + 7

        def get_kmp_next(s):
            next_arr = [0] * len(s)
            i, j = 1, 0

            while i < len(s):
                while j > 0 and s[i] != s[j]:
                    j = next_arr[j - 1]

                if s[i] == s[j]:
                    j += 1

                next_arr[i] = j
                i += 1
            return next_arr

        kmp_next = get_kmp_next(evil)

        @lru_cache(None)
        def solve(idx, is_prefix_s1, is_prefix_s2, prefix_evil):
            if prefix_evil == len(evil):
                return 0
            if idx == n:
                return 1

            total = 0
            first = ord(s1[idx]) if is_prefix_s1 else ord('a')
            last = ord(s2[idx]) if is_prefix_s2 else ord('z')

            for c in range(first, last + 1):
                char = chr(c)

                next_prefix_s1 = is_prefix_s1 and c == first
                next_prefix_s2 = is_prefix_s2 and c == last

                next_prefix_evil = prefix_evil
                while next_prefix_evil and evil[next_prefix_evil] != char:
                    next_prefix_evil = kmp_next[next_prefix_evil - 1]

                if char == evil[next_prefix_evil]:
                    next_prefix_evil += 1

                total += solve(idx + 1, next_prefix_s1, next_prefix_s2, next_prefix_evil)
                total %= mod

            return total

        return solve(0, True, True, 0)