[코테] 리트코드 트라이 336. Palindrome Pairs

336. Palindrome Pairs

문제 링크

https://leetcode.com/problems/palindrome-pairs/description/

문제 설명

You are given a 0-indexed array of unique strings words.

A palindrome pair is a pair of integers (i, j) such that:

• 0 <= i, j < words.length,
• i != j, and
• words[i] + words[j] (the concatenation of the two strings) is a palindrome.

Return an array of all the palindrome pairs of words.

You must write an algorithm with O(sum of words[i].length) runtime complexity.

제한 사항

• 1 <= words.length <= 5000
• 0 <= words[i].length <= 300
• words[i] consists of lowercase English letters.

입출력 예 #1

1. Input : words = [“abcd”,”dcba”,”lls”,”s”,”sssll”]
2. Output : [[0,1],[1,0],[3,2],[2,4]]
3. Explanation : The palindromes are [“abcddcba”,”dcbaabcd”,”slls”,”llssssll”]

입출력 예 #2

1. Input : words = [“bat”,”tab”,”cat”]
2. Output : [[0,1],[1,0]]
3. Explanation : The palindromes are [“battab”,”tabbat”]

입출력 예 #3

1. Input : words = [“a”,””]
2. Output : [[0,1],[1,0]]
3. Explanation : The palindromes are [“a”,”a”]

문제 풀이

search는 해당 단어가 있는지 여부를 확인하기 위해 마지막 문자의 값을 True로 해주고, startswith는 search와 동일하나 문자열의 문자가 있는지를 순회하며 확인한다.

# 브루트포스 풀이 : 시간초과
class Solution:
    def palindromePairs(self, words: List[str]) -> List[List[int]]:
        def is_palindrome(word):
            return word == word[::-1]

        output = []
        for i, word1 in enumerate(words):
            for j, word2 in enumerate(words):
                if i == j:
                    continue
                if is_palindrome(word1 + word2):
                    output.append([i, j])
        
        return output

# 트라이를 저장할 노드
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)):
            if self.is_palindrome(word[0:len(word) - i]):
                node.palindrome_word_ids.append(index)
            node = node.children[char]
            node.val = char
        node.word_id = index

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

        while word:
            # 판별 로직 1. 탐색 중간에 word_id가 있고, 나머지 문자가 팰린드롬인 경우
            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:]
        
        # 판별 로직 2. 끝까지 탐색했을 때 word_id가 있는 경우
        if node.word_id >= 0 and node.word_id != index:
            result.append([index, node.word_id])

        # 판별 로직 3. 끝까지 탐색했을 때 palindrome_word_ids가 있는 경우 
        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