백준[17610] 양팔저울 (Python3)

추를 통해 잴 수 없는 모든 경우의 수를 세는 문제

처음에 시간 복잡도 계산을 잘못해서 맞는 풀이임에도 계속 고민하다 생각보다 시간을 오래 썼다.
애매한 지식이 실수를 만들었다.

완전 탐색

3가지 경우에 대해 생각할 수 있다.

1. 추를 놓지 않는 경우
2. 추를 올려놓는 경우
3. 추를 반대쪽에 올려놓는 경우

2번의 경우 현재 무게에서 +를 해주고, 3번의 경우 -를 해준다.
이렇게 나오는 각 무게들에 방문 표시를 해준 뒤
1부터 모든 추의 합까지 탐색하면서 방문 처리되지 않은 값을 세면 된다.

 

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()