[Programmers/Lv. 2/Python] 바이러스 파이프

📎 문제 정보

난이도	Lv. 2 / 2025 카카오 하반기 1차
유형	완전탐색, DFS/그래프 탐색, 트리
언어	Python
플랫폼 / 제목	Programmers / 바이러스 파이프

---

🔍 문제 분석

• n개의 배양체가 n-1개의 파이프로 트리 형태로 연결되어 있고, 하나의 배양체가 처음부터 감염되어 있다.
• 파이프는 종류 단위로만 열고 닫을 수 있고, 열려 있는 동안 감염이 인접 노드로 전파된다. 

💡 풀이 아이디어

• 파이프 타입별로 인접 리스트 arr[node][type]를 만들어, 특정 타입을 열었을 때 해당 타입 간선만 따라 감염이 퍼지도록 구성한다.
• k번의 행동을 (A, B, C)의 길이 k 중복순열로 모두 나열하고 product([1,2,3], repeat=k), 각 순열대로 순서대로 파이프를 열어 감염 상태를 갱신한 뒤 감염 수의 최댓값을 추적한다.

---

💻 코드

from itertools import product


def infect(arr, opened, start, pipe):
    if not arr[start][pipe]:
        return

    for next_n in arr[start][pipe]:
        if opened[next_n] == 1: continue
        opened[next_n] = 1
        infect(arr, opened, next_n, pipe)

    return


def solution(n, infection, edges, k):
    arr = [[[] for _ in range(4)] for _ in range(n + 1)]

    for x, y, c_type in edges:
        arr[x][c_type].append(y)
        arr[y][c_type].append(x)

    pipes_candidates = product([1, 2, 3], repeat=k)
    answer = 1

    for pipes_candidate in pipes_candidates:
        opened = [0] * (n + 1)
        opened[infection] = 1

        for pipe in pipes_candidate:
            for i in range(1, n + 1):
                if opened[i] == 1:
                    infect(arr, opened, i, pipe)

        cnt = sum(opened)
        answer = max(answer, cnt)

    return answer

 

---

⏱️ 시간/공간 복잡도

시간 복잡도	O(3^k · k · n) - 후보 개수 3^k, k번 파이프 개방 × 매번 노드 최대 n개 순회 + DFS 전파
공간 복잡도	O(n)

 

---

📝 배운 점 / 실수했던 부분

• itertools.permutations(iterable, r)은 중복 순열이 아니다! product를 써야한다!!! 
• 현재 내 코드는.. 매우 비효율적이다.
  • 검색을 해보면, 비트마스크, 큐 BFS, 메모이제이션 등 다양한 해결 방법이 있어보인다..
  • product로 전체 후보를 만드는 것이 아니라, DFS/백트래킹 구조로 체크를 해야한다.
  • 각 단계에서 갈 수 있는 방향으로 하나씩 끝까지 가보고, 돌아올 때마다 지금까지 본 것 중 최댓값을 챙겨서 위로 넘긴다.
  • 이렇게 한다면 시간 복잡도는 O(n · 2^k)가 된다!

from collections import deque
from itertools import product


def infect(arr, opened, pipe, n):
    q = deque(i for i in range(1, n + 1) if opened[i])
    while q:
        cur = q.popleft()
        for next_n in arr[cur][pipe]:
            if not opened[next_n]:
                opened[next_n] = 1
                q.append(next_n)


def solution(n, infection, edges, k):
    arr = [[[] for _ in range(4)] for _ in range(n + 1)]
    for x, y, c_type in edges:
        arr[x][c_type].append(y)
        arr[y][c_type].append(x)

    def dfs(depth, last_pipe, opened):
        cnt = sum(opened)
        best = cnt                      

        if cnt == n or depth == k:
            return best

        for pipe in (1, 2, 3):
            if pipe == last_pipe:
                continue
            new_opened = opened[:]
            infect(arr, new_opened, pipe, n)
            best = max(best, dfs(depth + 1, pipe, new_opened))

        return best

    opened = [0] * (n + 1)
    opened[infection] = 1
    return dfs(0, 0, opened)

pipe=1 박스는 두 자식들 중에서 큰 값을 받아서, depth가 0인 박스로 올려보낸다.