[프로그래머스] 가사 검색 (파이썬)

🗂️   문제 

2020 KAKAO BLIND RECRUITMENT_가사 검색

---

 
 
 

---

📌   Point

트라이(Trie) 자료구조

• 단어의 길이별로 Trie 자료구조를 만들어서 효율적으로 쿼리를 처리할 수 있도록 한다.
• 정방향 Trie와 역방향 Trie를 동시에 초기화하고 단어를 삽입한다.
  • 정방향 Trie : 접미사에 "?"가 있을 때 사용
  • 역방향 Trie : 접두사에 "?"가 있을 때 사용

---

📁  데이터 구조 예시

 

📎  입력 단어

words = ["apple", "app", "bat"]

 

📎  길이별 Trie 초기화

# 정방향 Trie (tries)
tries = {
    3: Trie(),  # 3글자 단어용 Trie
    5: Trie()   # 5글자 단어용 Trie
}

# 역방향 Trie (tries_reverse)
tries_reverse = {
    3: Trie(),
    5: Trie()
}

 

📎  정방향 Trie 구성
 

• tries[3]에 "app", "bat" 삽입.
• tries[5]에 "apple" 삽입.

    Trie 구조 (정방향, 길이 3)

        (root)
       /    \
      a      b
      |      |
      p      a
      |      |
      p      t

 
   
    Trie 구조 (정방향, 길이 5)

         (root)
            |
            a
            |
            p
            |
            p
            |
            l
            |
            e

 

📎  역방향 Trie 구성
 

• tries_reverse[3]에 "ppa", "tab" 삽입.
• tries_reverse[5]에 "elppa" 삽입.

    Trie 구조 (정방향, 길이 3)

        (root)
       /    \
      p      t
      |      |
      p      a
      |      |
      a      b

 
    Trie 구조 (정방향, 길이 5)

         (root)
            |
            e
            |
            l
            |
            p
            |
            p
            |
            a

 
 
 
 

---

📄   코드

class TrieNode:
    def __init__(self):
        self.children = {}  # 자식 노드를 저장할 딕셔너리
        self.count = 0      # 현재 노드를 경유하는 단어의 수

class Trie:
    def __init__(self):
        self.root = TrieNode()  # 루트 노드 초기화
    
    def insert(self, word):
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
            node.count += 1
    
    def count_match(self, query):
        # 전체가 "?"인 경우, 루트의 count를 반환한다.
        if set(query) == {'?'}:
            return self.root.count
        
        node = self.root
        for char in query:
            if char == '?':
                break
            if char not in node.children:
                return 0
            node = node.children[char]
        
        return node.count if node else 0

def solution(words, queries):
    tries = {}         # 정방향 Trie
    tries_reverse = {} # 역방향 Trie
    length_words = {}  # 길이별 단어 개수 저장. {'길이' : 개수}
    
    # 단어들을 정방향 및 역방향 Trie에 삽입
    for word in words:
        length = len(word)
        
        # 길이별 단어 개수 저장
        length_words[length] = length_words.get(length, 0) + 1
        
        if length not in tries:
            tries[length] = Trie()
            tries_reverse[length] = Trie()
        
        tries[length].insert(word)
        tries_reverse[length].insert(word[::-1])
    
    # 각 쿼리에 대해 매칭되는 단어 수 계산
    result = []
    for query in queries:
        length = len(query)
        
        # 전부 "?"로 이루어진 경우
        if set(query) == {'?'}:
            result.append(length_words.get(length, 0))
            continue
        
        # 접두사에 "?"가 있는 경우
        if query[0] == '?':
            if length not in tries_reverse:
                result.append(0)
                continue
            
            reversed_query = query[::-1]
            match_count = tries_reverse[length].count_match(reversed_query)
            result.append(match_count)
        
        # 접미사에 "?"가 있는 경우
        else:
            if length not in tries:
                result.append(0)
                continue
            
            match_count = tries[length].count_match(query)
            result.append(match_count)
    
    return result

 
 
 
 

---

✍🏻   풀이

시간 복잡도 : O(NM + QM)

N: 단어수,  M: 단어 길이,  Q: 쿼리 수

 

1. 클래스 구조
   • Class : TrieNode
     • Node
       • children : 다음 문자로 연결되는 자식 노드를 저장하는 딕셔너리
       • count     : 해당 노드까지 오는 단어의 개수를 저장
   • Class : Trie
     • 트라이의 root 노드를 초기화한다.
     • insert(word)  :  단어를 트라이에 삽입한다.
       • 각 문자마다 노드를 생성하고 해당 경로의 count를 1씩 증가한다.
     • count_match(query)  :  쿼리와 매칭되는 단어의 수를 반환한다.
       • query가 "?"로만 이루어진 경우, 해당 길이의 단어 개수를 바로 반환
       • "?" 문자가 나오면 query 문자 탐색을 중단  →   노드의 count를 반환하여 단어 수 계산
2. 단어 처리
   • 길이별로 단어를 처리
   • 단어의 길이를 키(length)로 하여 Trie를 동적으로 생성한다.
     • tries : 정방향 Trie
     • tries_reverse : 역방향 Trie
   • Trie에 단어 삽입
     • (정방향) 문자열 그대로
     • (역방향) 문자열 뒤집어서
3. 쿼리 처리
   • 전부 "?"인 경우
     • length_words에서 해당 길이의 단어 개수를 바로 반환한다.
   • "?"가 앞에 붙는 경우 (접두사)
     • 역방향 트라이(tries_reverse) 사용
     • 쿼리를 뒤집어 "?"가 뒤에 위치하도록 한다.
       • 예를 들어 "??odo"인 경우, "odo??"로 바꾼다.
     • 접두사 부분까지 트라이를 따라간 후 해당 노드의 count를 반환한다.
   • "?"가 뒤에 붙는 경우 (접미사)
     • 정방향 트라이(tries) 사용
     • 접두사 부분까지 트라이를 따라간 후 해당 노드의 count를 반환한다.
       • 예를 들어 "fro??"인 경우, '?' 전인 'o'의 count를 반환한다.