[medium] 1268. Search Suggestions System

문제

https://leetcode.com/problems/search-suggestions-system/description/

---

풀이방법

찾고자 하는 단어(searchWord)랑 일치하는 알파벳 개수별 각각의 product들을 리스트로 출력하는 문제이다
1. 문자열 순서대로 출력해야하기 때문에 products 소팅
2. products로 trie 생성. 이 때 children과 함께 현재 노드에서 가능한 product들을 list로 관리
3. searchWord들의 알파벳을 trie에서 탐색하면서 해당 노드에서 가능한 product들을 최대 3개까지 list에 저장.
searchWord의 알파벳을 현재 node부터는 더이상 탐색이 불가능할 경우 flag를 False로 변경하고 빈 list 저장

반응형

---

코드

class TrieNode:
    def __init__(self):
        self.children = {}
        self.words = []

class Solution:
    def suggestedProducts(self, products: List[str], searchWord: str) -> List[List[str]]:
        ret = []
        self.root = TrieNode()
        flag = True #trie 탐색 계속할지 체크
        
        products.sort()
        
        #create trie
        for prod in products:
            node = self.root
            for c in prod:
                if c not in node.children:
                    node.children[c] = TrieNode()
                node = node.children[c]
                node.words.append(prod)

        node = self.root
        for c in searchWord:
            if flag and c in node.children:
                node = node.children[c]
                if len(node.words) > 3:
                    cur_ret = node.words[:3]
                else:
                    cur_ret = node.words
            else:
                cur_ret = []
                flag = False
            ret.append(cur_ret)
        return ret