백준[16934] 게임 닉네임 (Python3)

제한 시간 2초 치고는 데이터 범위가 작다 싶었는데,
키가 문자열이라 그런지 생각보다 실행 시간이 길었다. 트라이를 사용해 풀었다.

 

생각이 필요했던 부분이 입력된 입력된 닉네임이 고유한 닉네임이면 접두사를 별칭으로 해주고,
다음에 같은 닉네임을 입력받으면, 닉네임 + 중복 수를 별칭으로 해줘야 하는 거였다.

 

따라서, 삽입 전에 탐색을 통해 중복되는 닉네임이 있는지 확인할 필요가 있었다.
탐색에선 먼저 접두사를 구한 다음, 구해진 접두사의 길이가 닉네임의 길이와 같다면
중복된 닉네임으로 간주해 해당 닉네임의 카운트를 붙인 별칭을 반환해 출력했다.

 

코드

from collections import defaultdict

class Trie:
    def __init__(self):
        self.root = {}
        self.name = defaultdict(lambda: 1)

    def insert(self, S):
        current_node = self.root
        self.name[S] += 1
        for s in S:
            if s not in current_node:
                current_node[s] = {}
            current_node = current_node[s]
        current_node[-1] = True

    def search(self, S):
        current_node = self.root
        ret = []
        for s in S:
            ret.append(s)
            if s not in current_node:
                break
            current_node = current_node[s]
        nameCount = self.name[S]
        if len(S) == len(ret) and nameCount > 1:
            ret.append(str(nameCount))
        return "".join(ret)

def main():
    trie = Trie()
    for i in range(int(input())):
        S = input()
        print(trie.search(S))
        trie.insert(S)

if __name__ == "__main__":
    main()