336. Palindrome Pairs

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

Python

Remember Suffix in Trie Tree

K is the length of logest word; N is length of words

• Time: O(K^2+N)
• Space: O((K+N)^2)

Consider 3 cases in Word1 and Word2:

1. Word1 is reversed Word2: CATTAC
2. Word2 suffix is reversed Word1: CATSOLOSTAC
3. Word1 prefix is reversed Word2: CATSOLOSTAC

class Solution:
    def palindromePairs(self, words: List[str]) -> List[List[int]]:
        # Build the trie tree
        trie = {}
        for i, word in enumerate(words):
            cur = trie
            word = word[::-1] # Build the trie tree with inverse order of the word
            for j, letter in enumerate(word):
                if word[j:] == word[j:][::-1]:
                    if 'suffix' not in cur:
                        cur['suffix'] = []
                    cur['suffix'].append(i)   # Remember any suffix id (index) on the level if exist
                cur = cur.setdefault(letter, {})
            cur['eow'] = i

        # Consider cases and find solutions
        solutions = []
        for i, word in enumerate(words):
            cur = trie
            for j, letter in enumerate(word):
                # Case 3
                if cur.get('eow') is not None:
                    if word[j:] == word[j:][::-1]:
                        solutions.append([i, cur.get('eow')])
                if letter not in cur:
                    break
                cur = cur[letter]
            else:
                # Case 1
                if cur.get('eow') is not None and cur['eow'] != i:
                    solutions.append([i, cur['eow']])

                # Case 2
                for j in cur.get('suffix', []):
                    solutions.append([i, j])

        return solutions