[백준] 14725 - 개미굴 [Python(파이썬)]

문제

www.acmicpc.net/problem/14725

14725번: 개미굴

 

풀이

먹이(단어)를 하나의 노드로 간략하게 트라이(Trie)를 구현한다. 처음엔 스택을 이용해서 탐색하려 했으나 백트래킹하는 과정에서 재귀를 사용하는 것이 더 나을 것 같아서 변경했다. 이해를 돕기 위해 아래에 트라이 구현 코드까지 올렸다.

 

코드

문제풀이

import sys

class Trie:
    def __init__(self):
        self.root = {} # children

    def add(self, foods):
        cur = self.root
    
        for food in foods:
            if food not in cur:
                cur[food] = {} # dictonary => children
            cur = cur[food]
        cur[0] = True # leaf node
    
    def travel(self, level, cur):
        if 0 in cur:
            return

        cur_children = sorted(cur)

        for child in cur_children:
            print("--"*level + child)
            self.travel(level+1, cur[child])      

N = int(sys.stdin.readline().strip())
nest = Trie()
for _ in range(N):
    input_line = sys.stdin.readline().split()
    K = input_line[0 ]
    foods = input_line[1:]
    nest.add(foods)
nest.travel(0, nest.root)

 

 

트라이

class Node:
    def __init__(self):
        self.word = False
        self.children = {}
    
class Trie:
    def __init__(self):
        self.root = Node()
    
    def insert(self, word: str) -> None:
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = Node()
            node = node.children[char]
        node.word = True
    
    def search(self, word: str) -> bool:
        node = self.root
        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]

        return node.word
    
    def startsWith(self, prefix: str) -> bool:
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        
        return True
    
    def travel(self, level: int, node: Node) -> None: 
        if node.word:
            return

        cur_children = sorted(node.children)

        for child in cur_children:
            print("--"*level + child)
            self.travel(level+1, node.children[child])