백준[16906] 욱제어 (Python3)

dfs만 써도 풀 수 있지만, 연습을 위해 트라이를 사용했다
예제 볼 때는 단어 길이가 오름차순인줄 알았지만, 제출하고 보니 무작위였다.

트라이 풀이

매 탐색마다 0 또는 1을 붙이는 2가지 분기로 나뉜다.
만약, 현재 만들어진 단어가 다른 단어의 접두사인 경우(트라이에 끝이 명시돼 있음)

탐색을 중단하고 다른 단어를 알아봐야 한다.
이렇게 생긴 문자열이 만들려는 단어 길이와 일치하면 해당 단어를 트라이에 추가해준다.

 

import sys
sys.setrecursionlimit(10 ** 5)
class Trie:
    def __init__(self):
        self.root = {}
        self.answer= []

    def insert(self, word):
        current_node = self.root
        for w in word:
            if w not in current_node:
                current_node[w] = {}
            current_node = current_node[w]
        current_node[-1] = True

    def search(self, word):
        current_node = self.root
        for w in word:
            if -1 in current_node:
                return False
            if w not in current_node:
                return True
            current_node = current_node[w]
        if -1 in current_node:
            return False
        return True

def setting():
    N = int(input())
    lens = list(map(int, input().split()))
    words = []
    for i in range(N):
        words.append((lens[i], i))
    words = sorted(words, key=lambda x: x[0])
    return words

def dfs(trie, length, word, index):
    if len(word) >= 1 and not trie.search(word):
        return
    if len(word) == length:
        trie.insert(word)
        trie.answer[index] = "".join(list(map(str, word)))
        return True
    else:
        if dfs(trie, length, word + [0], index):
            return True
        if dfs(trie, length, word + [1], index):
            return True
    return

def main():
    trie = Trie()
    words = setting()
    trie.answer = [-1] * len(words)
    for length, i in words:
        dfs(trie, length, [], i)
    if -1 in trie.answer:
        print(-1)
    else:
        print(1)
        for a in trie.answer:
            print(a)

if __name__ == "__main__":
    main()