문제 요약
N개의 배양체가 N - 1개의 파이프로 연결되어 하나의 트리를 이룹니다.
각 파이프의 종류는 A, B, C 중 하나이며, 처음에는 모든 파이프가 닫혀 있습니다. 한 번의 행동으로 특정 종류의 파이프를 모두 열었다가 닫을 수 있습니다.
파이프가 열려 있는 동안에는 이미 감염된 배양체로부터 열린 파이프를 통해 연결된 배양체로 바이러스가 퍼집니다.
최대 K번 행동했을 때 감염시킬 수 있는 배양체 수의 최댓값을 구해야 합니다.
핵심 아이디어
1. 같은 타입의 연결 요소 전체가 한 번에 감염된다
예를 들어 A 타입 파이프를 열었다고 가정해 보겠습니다.
현재 감염된 배양체와 A 타입 파이프만 사용해 연결된 모든 배양체가 감염됩니다. 새롭게 감염된 배양체에서도 다시 A 타입 파이프를 따라 바이러스가 퍼질 수 있습니다.
따라서 한 타입의 파이프를 열었을 때의 결과는 다음과 같이 생각할 수 있습니다.
해당 타입의 파이프만 남겨 만든 그래프에서, 감염된 노드가 하나라도 포함된 연결 요소 전체가 감염된다.
이 연결 요소는 트리와 파이프 정보가 바뀌지 않는 한 항상 동일합니다. 따라서 A, B, C 타입별 연결 요소를 미리 구할 수 있습니다.
2. 감염 상태를 비트마스크로 표현한다
배양체는 최대 100개입니다. Python의 정수는 크기 제한 없이 비트를 사용할 수 있으므로 감염 상태를 하나의 정수로 나타낼 수 있습니다.
i번 비트가 1: i + 1번 배양체가 감염됨
i번 비트가 0: i + 1번 배양체가 감염되지 않음예를 들어 1번, 3번, 4번 배양체가 감염되었다면 다음과 같이 표현할 수 있습니다.
1101₂연결 요소 역시 비트마스크로 저장합니다. 현재 감염 상태와 연결 요소의 비트 AND 연산 결과가 0이 아니라면, 그 연결 요소에는 감염된 배양체가 하나 이상 존재한다는 뜻입니다.
if infected & component:
infected |= component이때 OR 연산을 사용하면 연결 요소에 속한 모든 배양체를 감염 상태에 추가할 수 있습니다.
3. 가능한 감염 상태를 반복 DP로 관리한다
각 행동에서 선택할 수 있는 파이프 종류는 3개이고, 행동 횟수는 최대 10번입니다.
가능한 선택 순서의 수는 최대 다음과 같습니다.
3^10 = 59,049충분히 탐색할 수 있는 크기입니다.
모든 행동 순서를 그대로 저장할 필요는 없습니다. 서로 다른 순서로 파이프를 열었더라도 같은 횟수에 동일한 감염 상태에 도달했다면, 이후 가능한 결과도 완전히 같기 때문입니다.
states = 현재 행동 횟수까지 도달할 수 있는 감염 상태들의 집합각 행동에서는 states의 모든 감염 상태에 대해 A, B, C 타입을 각각 열어 보고, 그 결과를 next_states에 넣습니다.
Python의 set은 같은 정수를 중복해서 저장하지 않습니다. 감염 상태를 비트마스크 정수로 표현했으므로, 같은 상태에 여러 번 도달해도 자동으로 하나만 남습니다.
따라서 재귀 함수나 lru_cache 없이도 중복 상태를 제거할 수 있습니다.
최대 K번과 정확히 K번
문제에서는 행동을 최대 K번 할 수 있다고 했습니다.
파이프를 열어도 기존에 감염된 배양체가 회복되지는 않으므로 감염 수는 절대 감소하지 않습니다. 효과가 없는 타입을 다시 열더라도 현재 상태가 그대로 유지될 뿐입니다.
따라서 정확히 K번 행동하는 경우를 탐색해도 최대 K번 행동했을 때의 최댓값과 같습니다.
알고리즘
- 파이프 타입별 인접 리스트를 만듭니다.
- 각 타입의 파이프만 이용해 연결 요소를 구합니다.
- 각 연결 요소를 비트마스크로 저장합니다.
- 최초 감염 배양체만 포함하는 비트마스크를 만듭니다.
- 현재 가능한 감염 상태들을
states집합에 저장합니다. - 각 상태에서 파이프 타입
A,B,C를 여는 경우를 모두 계산합니다. - 새로 만들어진 감염 상태를
next_states집합에 저장하여 중복을 제거합니다. - 이 과정을
K번 반복합니다. - 최종 상태 중 1비트의 개수가 가장 많은 값을 반환합니다.
Python 풀이
def solution(n, infection, edges, k):
# 타입별 인접 리스트
graph = [[[] for _ in range(n)] for _ in range(3)]
for x, y, pipe_type in edges:
x -= 1
y -= 1
pipe_type -= 1
graph[pipe_type][x].append(y)
graph[pipe_type][y].append(x)
# 각 타입의 연결 요소를 비트마스크로 저장
components = [[] for _ in range(3)]
for pipe_type in range(3):
visited = [False] * n
for start in range(n):
if visited[start]:
continue
stack = [start]
visited[start] = True
component = 0
while stack:
node = stack.pop()
component |= 1 << node
for next_node in graph[pipe_type][node]:
if not visited[next_node]:
visited[next_node] = True
stack.append(next_node)
components[pipe_type].append(component)
def spread(infected, pipe_type):
result = infected
for component in components[pipe_type]:
# 감염된 노드가 하나라도 있는 연결 요소 전체를 감염시킨다.
if infected & component:
result |= component
return result
initial_infected = 1 << (infection - 1)
states = {initial_infected}
for _ in range(k):
next_states = set()
for infected in states:
for pipe_type in range(3):
next_states.add(spread(infected, pipe_type))
states = next_states
return max(infected.bit_count() for infected in states)예제 1 분석
n = 10
infection = 1
edges = [
[1, 2, 1],
[1, 3, 1],
[1, 4, 3],
[1, 5, 2],
[5, 6, 1],
[5, 7, 1],
[2, 8, 3],
[2, 9, 2],
[9, 10, 1],
]
k = 2먼저 B 타입 파이프를 열면 1번과 연결된 5번 배양체가 감염됩니다.
감염 상태: {1, 5}이후 A 타입 파이프를 열면 다음과 같이 감염이 퍼집니다.
1번 -> 2번, 3번
5번 -> 6번, 7번최종 감염 상태는 다음과 같습니다.
{1, 2, 3, 5, 6, 7}따라서 감염된 배양체의 최대 개수는 6입니다.
왜 연결 요소를 미리 구할 수 있을까?
한 타입의 파이프를 여는 동안 감염은 해당 타입의 파이프를 따라 더 이상 퍼질 곳이 없을 때까지 계속됩니다.
현재 감염된 노드와 같은 연결 요소에 속한 노드는 모두 감염되고, 다른 연결 요소로는 이동할 수 없습니다.
파이프의 종류와 연결 관계는 행동 도중 바뀌지 않으므로 타입별 연결 요소를 매번 다시 탐색할 필요가 없습니다. 처음 한 번만 구한 뒤 모든 DP 상태에서 재사용할 수 있습니다.
복잡도 분석
한 행동 횟수에서 관리하는 서로 다른 감염 상태 수를 S라고 하겠습니다.
행동마다 세 타입을 선택할 수 있으므로 상태 수의 상한은 다음과 같습니다.
S = O(3^K)각 상태에서 세 타입을 시도하며, 한 번의 확산에서는 해당 타입의 연결 요소를 최대 N개 확인합니다.
- 전처리 시간 복잡도:
O(N) - 탐색 시간 복잡도:
O(3^K × N) - 공간 복잡도:
O(3^K + N)
set을 사용하므로 실제 상태 수는 중복이 제거되어 이 상한보다 작을 수 있습니다. 제한에서 N ≤ 100, K ≤ 10이므로 충분히 처리할 수 있습니다.
정리
이 문제의 핵심은 파이프를 하나씩 따라가며 감염을 시뮬레이션하는 대신, 같은 타입 파이프로 이루어진 연결 요소 단위로 감염을 처리하는 것입니다.
또한 Python의 정수를 비트마스크로 사용하면 최대 100개 배양체의 감염 상태를 간결하게 저장할 수 있습니다.
타입별 연결 요소 전처리
→ 감염 상태를 비트마스크로 표현
→ 세 가지 행동을 반복 DP로 탐색
→ set으로 같은 감염 상태를 제거파이프 타입이 3개, 최대 행동 횟수가 10번이라는 작은 선택지를 활용하면 모든 가능한 행동 순서를 효율적으로 비교할 수 있습니다.
