[Leetcode] 리트코드 - 펠린드롬 페어(palindrome-pairs) 파이썬(python) 풀이

문제

단어 리스트에서 words[i] + words[j]가 펠린드롬이 되는 모든 인덱스 조합(i, j)를 구하라.

입력

["abcd","dcba","lls","s","sssll"]

출력

[[0,1],[1,0],[3,2],[2,4]]

["dcbaabcd","abcddcba","slls","llssssll"]이 팰린드롬이다.

입력

["bat","tab","cat"]

출력

[[0,1],[1,0]]

["battab","tabbat"]이 팰린드롬이다.

코드

import collections
from typing import List


# 트라이 저장용 노드.
class TrieNode:
    def __init__(self):
        self.children = collections.defaultdict(TrieNode)
        self.word_id = -1
        self.palindrome_word_ids = []


class Trie:
    def __init__(self):
        self.root = TrieNode()

    # 펠린드롬 판별 함수.
    @staticmethod
    def is_palindrome(word: str) -> bool:
        return word[::] == word[::-1]

    # 단어 삽입.
    def insert(self, index, word) -> None:
        node = self.root
        for i, char in enumerate(reversed(word)):
            # 두번째 로직에서 문자를 하나씩 제거해가며 팰린드롬을 판별해서 체크함, 여기 word 리스트를 잘라서 체크함 wordreverse 아님!!!!!!!
            if self.is_palindrome(word[0:len(word) - i]):
                node.palindrome_word_ids.append(index)
            node = node.children[char]
        node.word_id = index

    def search(self, index, word) -> List[List[int]]:
        result = []
        node = self.root

        # 여기 while 에서 node = nod.children[word[0]]으로 node 끝까지 타고 가서 밑에 판별 로직 1, 2에서 끝점이랑 비교가 가능해짐.
        while word:
            # 판별 로직 3
            if node.word_id >= 0:
                if self.is_palindrome(word):
                    result.append([index, node.word_id])
            if not word[0] in node.children:
                return result
            node = node.children[word[0]]
            word = word[1:]

        """
        판별 로직 1 
        위쪽 while 문에서 구현된 트라이와 word가 맞을 때 까지 쭉 안으로 들어오고 
        그 지점에서 if 문이 실행됨.
        즉, 꺼꾸로 들어간 트라이와 원래의 word 값이 같은 경우에
        밑의 if문으로 걸러짐. 
        """
        if node.word_id >= 0 and node.word_id != index:
            result.append([index, node.word_id])

        """
        판별 로직 2
        위쪽 while 문에서 구현된 트라이와 word가 맞을 때 까지 쭉 안으로 들어오고
        그 지점에서 for 문이 실행됨.
        그 자리에 palindrome_word_id가 있으면 그것도 정답. 
        """
        for palindrome_word_id in node.palindrome_word_ids:
            result.append([index, palindrome_word_id])

        return result


class Solution:
    def palindromePairs(self, words: List[str]) -> List[List[int]]:
        trie = Trie()

        for i, word in enumerate(words):
            trie.insert(i, word)

        results = []
        for i, word in enumerate(words):
            results.extend(trie.search(i, word))

        return results


a = Solution().palindromePairs(['d', 'cbbcd', 'dcbb', 'dcbc', 'cbbc', 'bbcd'])
print(a)

풀이

트라이로 O(n)으로 풀이한다. 팰린드롬을 판별해야 하기 때문에, 뒤집어서 트라이로 구성하면 해법을 찾을 수 있을 것 같다.

먼저, 이 예제 입력값은 너무 단순하므로, ['d', 'cbbcd', 'dcbb', 'dcbc', 'cbbc', 'bbcd'] 정도로 좀 더 복잡한 값을 입력값으로 해서 뒤집어서 밑의 그림과 같이 트라이로 구현하자.

(위 그림에서 왼쪽트리에 dcbdc가 아니라 dcbbc임.)

입력값을 뒤집으면 ['d', 'dcbbc', 'bbcd', 'cbcd', 'cbbc', 'dcbb']가 된다. 뒤집은 다음, 문자 단위로 계속 내려가면서 트라이를 구현하고, 각각의 단어가 끝나는 지점에는 단어 인덱스를 word_id로 부여했다. 코드에서는 word_id이고, 위의 그림에서는 w로 표시했다.

첫 번째 로직

단어를 뒤집어서 구축한 트라이 이기 때문에 입력값은 순서대로 탐색하다가, 끝나는 지점의 word_id 값이 -1아니라면, 현재 인덱스 index와 해당 word_id는 팰린드롬으로 판단할 수 있다. 예를 들어 위의 그림에서 입력값 bbcd의 트라이 탐색이 끝나는 지점에서 word_id = 2가 셋팅되어 있고, bbcd의 인덱스는 5이기 때문에, [5, 2]인 bbcd + dcbb는 팰린드롬이며, 이는 정답 중 하나가 된다. 여기서 첫 번째 로직으로 [2, 5], [5, 2]가 걸러진다.

두 번째 로직

트라이 삽입 중에 남아 있는 단어가 팰린드롬이라면 미리 팰린 드롬 여부를 세팅해 두는 방법이다. 즉 입력값 ['d', 'cbbcd', 'dcbb', 'dcbc', 'cbbc', 'bbcd']에서 cbbc는 단어 자체가 팰린드롬이므로 루트에 바로 입력값의 인덱스인 p = 4를 셋팅하고, word[0:len(word) - i] 형태로 단어에서 문자 수를 계속 줄여 나가며 팰린드롬 여부를 체크한다. 문자가 하나만 남게 될 때는 항상 팰린드롬이므로 마찬가지로 p = 4를 마지막에 셋팅한다. 당연히 이 마지막 값은 항상 w의 바로 앞 노드가 된다.

참고로 위의 그림에서 p로 표현한 것을 코드에서는 palindrome_word_ids로 풀어서 표현했다. 또한 코드에서는 속성의 이름을 복수형으로 정했는데, 그 이유는 이 그림의 루트 경우처럼 p 값이 여러 개가 될 수 있기 때문이다. 따라서 코드에서 palindrome_word_ids는 리스트이며 복수형이다.

이제 남아 있는 단어가 팰린드롬인 경우를 좀 더 살펴보자. 위의 그림에서 w는 단어의 끝이고 p는 w이전 노드에 반드시 셋팅이 된다. 문자가 하나만 남았을 때는 항상 팰린드롬이기 때문이다. 다시 한번 입력 값 ['d', 'cbbcd', 'dcbb', 'dcbc', 'cbbc', 'bbcd']을 보면 dcbb의 인덱스는 2이고, 트라이에서는 d → c → b → b의 마지막 노드가 p = 1이다. 그렇다면 [2, 1]은 정답 중 하나가 된다. 실제로 인덱스 1의 입력 값은 cbbcd이므로 dcbb + cbbcd는 팰린드롬으로, 정답 중 하나다. 또 다른 경우로, 인덱스 0의 d를 살펴보자. 위의 그림에서 d 노드는 p = 1이다. 즉 [0, 1]도 정답이 된다. d + cbbcd이며 마찬가지로 팰린드롬이다(주석 참고.) 즉, 두 번째 로직으로는 [0, 1], [2, 1]이 걸러진다.

세 번째 로직

마지막인 세 번째 판별 로직은 입력값을 문자 단위로 확인해 나가다가 해당 노드의 word_id가 -1이 아닐 때, 나머지 문자가 팰린드롬이라면 팰린드롬으로 판별하는 경우다. dcbc + d가 이에 해당하는데, 입력값 dcbc는 먼저 d부터 탐색하다가 d의 word_id가 -1이 아니기 때문에 나머지 문자 cbc에 대한 팰린드롬 여부를 검사한다. 여기서는 dcbc + d를 팰린드롬으로 판별한다. 세 번째 로직에서는 [1, 4], [3, 0]이 걸러진다.

정리

3가지의 판별 로직을 다시 정리하면 다음과 같다.

1. 끝까지 탐색했을 때 word_id가 있는 경우.
2. 끝까지 탐색했을 때 palindrome_word_ids가 있는 경우.
3. 탐색 중간에 word_id가 있고 나머지 문자가 팰린드롬인 경우.

이렇게 3가지 경우를 팰린드롬으로 판별할 수 있으며, 입력값을 각각 한번씩만 대입하면 되기 때문에 O(n)으로 풀이할 수 있다. 좀 더 정확히는 단어의 최대 길이를 k로 했을 때 O(k^2n)이며, 브루트 포스로 풀이할 경우 O(kn^2)이다. 이 문제에서는 k가 훨씬 더 작기 때문에 트라이 풀이가 더 효율적이라고 볼 수 있다.

출처
파이썬 알고리즘 인터뷰 (글 : 박상길 그림 : 정진호) [책만]