[백준/파이썬][Platinum III] 전설 - 19585

[Platinum III] 전설 - 19585

문제 링크

성능 요약

메모리: 1661844 KB, 시간: 5816 ms

분류

자료 구조, 해시를 사용한 집합과 맵, 문자열, 트리, 트라이

제출 일자

2025년 2월 4일 00:45:12

문제 설명

Sogang ICPC Team에는 색상 이름과 닉네임의 순서로 이여서 팀명을 지으면 ICPC 리저널에서 수상할 수 있다는 전설이 있다. 색상 이름들과 닉네임들이 주어질 때, Q개의 팀에 대해 다음 리저널에서 수상할 수 있을지 전설에 기반해 알려주는 프로그램을 작성하자.

입력

첫째 줄에는 색상의 종류 C와 닉네임의 개수 N이 주어진다. (1 ≤ C, N ≤ 4,000)

다음 C개의 줄에는 색상 이름 C개가 한 줄에 하나씩 주어진다. 색상 이름은 중복되지 않는다.

다음 N개의 줄에는 Sogang ICPC Team 구성원들의 닉네임 N개가 한 줄에 하나씩 주어진다. 닉네임도 중복되지 않는다.

다음 줄에는 팀의 개수 Q가 주어진다. (1 ≤ Q ≤ 20,000)

다음 Q개의 줄에는 팀명 Q개가 한 줄에 하나씩 주어진다. 팀명은 중복되지 않는다.

모든 색상 이름의 길이와 닉네임의 길이는 1,000글자를 넘지 않는다. 모든 팀명은 2,000글자를 넘지 않는다. 모든 문자열은 알파벳 소문자로만 이루어져 있다.

출력

팀명이 입력된 순서대로, 전설에 따라 해당 팀이 다음 리저널에서 수상할 수 있다면

Yes

, 아니라면

No

를 출력한다.

 

풀이

아예 손도 못댄 문제. 

def check(word) :
    now = colors
    for i in range(len(word)) :
        if now.get(0) and word[i:] in names :
            return 1
        if not now.get(word[i]) :
            return
        now = now[word[i]]

C, N = map(int, input().split())
colors = {}
for _ in range(C) :
    now = colors
    for c in input().strip() :
        if not now.get(c) :
            now[c] = {}
        now = now[c]
    now[0] = 1

names = {input().strip() for i in range(N)}

for _ in range(int(input())) :
    print('Yes' if check(input().strip()) else 'No')

처음 참고한 풀이이다.

그리고 이를 활용해서 GPT 한테 어떤 자료구조와 알고리즘을 활용해서 풀이해달라고 요청하였다. 

class TrieNode:
    def __init__(self):
        self.children = {}
        self.is_end = False


class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str):
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.is_end = True

    def search_prefix(self, prefix: str):
        node = self.root
        for char in prefix:
            if char not in node.children:
                return None
            node = node.children[char]
        return node

    def exists(self, word: str):
        node = self.search_prefix(word)
        return node is not None and node.is_end


def check(word, color_trie, names):
    node = color_trie.root
    for i in range(len(word)):
        if node.is_end and word[i:] in names:  # 색상 부분이 존재하고, 나머지가 닉네임에 있으면
            return True
        if word[i] not in node.children:  # 색상 부분이 존재하지 않으면 종료
            return False
        node = node.children[word[i]]
    return False


def main():
    # 입력 받기
    C, N = map(int, input().split())

    color_trie = Trie()
    names = set()

    for _ in range(C):
        color_trie.insert(input().strip())

    for _ in range(N):
        names.add(input().strip())

    Q = int(input().strip())

    results = []
    for _ in range(Q):
        team_name = input().strip()
        results.append("Yes" if check(team_name, color_trie, names) else "No")

    # 결과 출력
    for res in results:
        print(res)


if __name__ == "__main__":
    main()

1. input().strip() : strip() 을 붙이지 않아서 오답을 여러번 내었다. 보통 다른 문제들은 이런 현상이 없었는데, 이 문제는 문자열 문제라 그런지 공백 관련해서도 엄청 빡빡하였다. 이제 붙이는 습관을 들여야 겠다.
2. 자료구조 공부 : 트라이라는 새로운 자료구조를 알게 되었고, 집합과 조합되어서 이런 문제도 해결하였다. 일반적인 코딩테스트에서는 트라이 까지는 안나오는 것으로 알고 있는데 이 문제를 내가 왜 풀고 있는지 의문이다. 
3. 위 코드는 다른 분의 코드를 그대로 가져왔다. 아직 트라이에 대해서 잘 몰라서 그러는지 봐도 잘 모르겠다.