두 단어의 조합으로 회문(Palindrome)을 만들 수 있는가(Leetcode - 336. Palindrome Pairs)

문제

풀이

브루트 포스

$O(N^2)$, TLE.

접두사, 접미사 이용하기

아이디어 : https://leetcode.com/problems/palindrome-pairs/discuss/79209/Accepted-Python-Solution-With-Explanation

핵심 아이디어는, 어떤 단어의 접두사가 회문이라면, 나머지 문자를 뒤집어 회문을 만들 수 있다는 것이다. 예를 들어 아래와 같은 세 단어가 있다고 하자.

words = ["bot", "t", "to"]

먼저 "bot"을 한 글자씩 분해해 보자. 그러면 "", "b", "bo", "bot" 4개 단어가 나온다. 이 중에서 회문은 공백을 포함해 "", "b" 2개다. 각각을 나머지 글자, 그리고 나머지를 뒤집은 글자, 그리고 원래 단어와 뒤집은 단어를 합쳐 나열해보면 다음과 같다.

뒤집은 접미사와 원래 단어를 결합한 단어는 회문이 됨을 알 수 있다. 뒤집은 접미사 중 "to"가 원래 단어 리스트에 들어 있으므로 "bot", "to"는 짝(pair)가 될 수 있다. 같은 방법으로 접두사에 대해서도 확인하면 모든 조합을 찾아낼 수 있다.

파이썬 구현은 다음과 같다.

class Solution:
    def palindromePairs(self, words: List[str]) -> List[List[int]]:
        def is_palindrome(string):
            return string == string[::-1]

        # map with index
        words = {word: i for i, word in enumerate(words)}

        palins = []
        for word, idx in words.items():
            for i in range(len(word) + 1):
                prefix = word[:i]
                suffix = word[i:]

                if is_palindrome(prefix):
                    rev = suffix[::-1]
                    if rev != word and rev in words:
                        palins.append([words[rev], idx])

                # i != len(word) must be checked,
                # because this means "" is being appended,
                # while "" is already prepended.
                if i != len(word) and is_palindrome(suffix):
                    rev = prefix[::-1]
                    if rev != word and rev in words:
                        palins.append([idx, words[rev]])

        return palins

Trie

이 문제의 출제 의도는 사실 Trie 자료구조를 사용하는 것이라고 생각한다. Palindrome + Trie 라고 생각하면 직관적으로 Trie with reversed words가 바로 떠오른다. 역순 트라이와 원래 단어를 비교해 회문을 만들 수 있는 경우는 다음 두 가지다. 편의상 찾고자 하는 단어를 A, 역순으로 뒤집은 트라이에 속하는 단어를 B라 하자.

1. B의 길이가 A의 길이보다 짧거나 같은 경우
2. B의 길이가 A의 길이보다 긴 경우

경우 1의 경우, 원래 단어의 문자를 역순 트라이의 노드값과 하나씩 비교해 나가면 된다. 그러면 원래 단어가 더 길어서 문자가 남는 경우가 발생하는데, 이 남는 문자가 회문인지 체크하면 된다.

경우 2의 경우도 마찬가지로 남는 노드의 문자로 회문을 만들 수 있는지 체크하면 된다. 다만 매번 노드에서 문자열을 만드는 것은 비효율적이므로, 트라이를 만들어 나가면서 남는 문자열이 회문인지를 체크하고, 회문이라면 회문이 시작하는 노드의 부모에 다음 문자열이 회문임을 표기한다.

class Trie:
    def __init__(self):
        self.children = collections.defaultdict(Trie)
        self.word_idx = -1
        self.palindrome_word_idx = []

    @staticmethod
    def is_palinderome(word):
        return word == word[::-1]

    def insert(self, word, index):
        node = self
        # make trie with reversed words
        for idx, char in enumerate(reversed(word)):
            # check if rest of word is a palindrome
            if self.is_palinderome(word[: len(word) - idx]):
                node.palindrome_word_idx.append(index)
            node = node.children[char]
        # mark index when traverse ends
        node.word_idx = index

    def search(self, word, index):
        res = []
        node = self

        while word:
            # Case 1. Check rest of the word is a palindrome
            if node.word_idx >= 0:
                if self.is_palinderome(word):
                    res.append(node.word_idx)

            if word[0] not in node.children:
                return res
            node = node.children[word[0]]
            word = word[1:]

        # Case 2. palindrome_word_idx already contains
        # every palindrome's indices.
        if node.word_idx >= 0:
            res.append(node.word_idx)
        res.extend(node.palindrome_word_idx)

        return res


class Solution:
    def palindromePairs(self, words: List[str]) -> List[List[int]]:
        node = Trie()
        for i, word in enumerate(words):
            node.insert(word, i)

        res = []
        for i, word in enumerate(words):
            res.extend([[i, cand] for cand in node.search(word, i) if i != cand])

        return res

벤치마크