[백준] (22860) 폴더 정리 (small) [Python]

문제 링크

https://www.acmicpc.net/problem/22860

문제

이름이 main 폴더 안에 여러가지 파일과 폴더가 존재한다.

main
 ├─ FolderA
 │    ├─ File1
 │    └─ File2
 └─ FolderB
       ├─ FolderC
       ├─ File1
       └─ File3

위 구조는 main 폴더의 하위 구조를 계층적으로 표시한 것이다. FolderA, FolderB, FolderC는 폴더이고 File1, File2, File3은 파일이다. 파일 이름이 같은 경우는 동일한 파일이다.

• 한 폴더 안에는 같은 이름의 파일이 두 개 이상 존재할 수 없다.
• main 하위 디렉토리에 같은 이름의 폴더가 두 개 이상 존재할 수 없다.

폴더 정리를 하기 위해 main 폴더 하위에 있는 파일들을 확인하려고 한다. 주어지는 쿼리에 대해 폴더와 파일의 정보를 알려주는 프로그램을 만들어보자.

---

코드

import sys
input = sys.stdin.readline
from collections import defaultdict, deque

class Node:
    def __init__(self, value, parent):
        self.value = value
        self.parent = parent
        self.child = list()
        self.files = set()
    def add_child(self, child):
        self.child.append(child)
    def remove_child(self, child):
        if child in self.child:
            self.child.remove(child)
    def add_file(self, file):
        self.files.add(file)
    def remove_file(self, file):
        if file in self.files:
            self.files.remove(file)

class Tree:
    def __init__(self, root: Node):
        self.root = root

n, m = map(int, input().strip().split())
root = Node("main", None)
tree = Tree(root)
visited = set()
visited.add("main")

inputs = deque()
for _ in range(n+m):
    inputs.append(input().strip())

while inputs:
    st = inputs.popleft()
    arr = list(st.split())
    if arr[0] not in visited:
        inputs.append(st)
        continue
    visited.add(arr[1])
    if int(arr[2]) == 1: # Folder
        stack = list()
        stack.append(tree.root)
        while True:
            cur = stack.pop()
            if cur.value == arr[0]:
                break
            for c in cur.child:
                stack.append(c)
        child = Node(arr[1], cur)
        cur.add_child(child)
    else: # File
        stack = list()
        stack.append(tree.root)
        while True:
            cur = stack.pop()
            if cur.value == arr[0]:
                break
            for c in cur.child:
                stack.append(c)
        cur.add_file(arr[1])

q = int(input().strip())
for _ in range(q):
    query = list(input().strip().split('/'))
    
    cur = tree.root
    for i in range(1, len(query)):
        for j in range(len(cur.child)):
            if cur.child[j].value == query[i]:
                break
        cur = cur.child[j]

    answer = defaultdict(int)
    for f in cur.files:
        answer[f] += 1
    for ch in cur.child:
        stack = list()
        stack.append(ch)
        while stack:
            curr = stack.pop()
            if curr.files:
                for f in curr.files:
                    answer[f] += 1
            if curr.child:
                for c in curr.child:
                    stack.append(c)

    print(len(answer), sum(answer[i] for i in answer))

 

처음에는 n+m의 입력이 제멋대로 주어지는 경우(main 하위 폴더 이전에 그 하위 폴더가 먼저 입력되는 경우 등)를 고려하지 못했기 때문에 IndexError가 발생했다. 그래서 visited라는 set을 추가해, 입력에서의 최상위 폴더가 visited에 등록되지 않은 경우는 큐의 맨 뒤로 이동하게 해서 문제를 해결했다.