21738(얼음깨기 펭귄) - 시간초과 - 해결(참고)

얼음깨기 펭귄

 

문제

도도는 심심해서 보드게임 카페에 갔다. 마침 평소에 즐겨 했던 얼음 깨기 펭귄의 업그레이드 버전으로 특수 얼음 깨기 펭귄 보드게임이 나와 직접 플레이해 보기로 결정했다. 특수 얼음 깨기 펭귄 게임은 특수 안경이 있어 특수 안경을 끼고 얼음들을 보면 얼음들 간의 연결 관계가 보인다.

특수 얼음 깨기 펭귄 게임에 있는 얼음의 종류로는 지지대의 역할을 하는 얼음과 일반 얼음 총 2가지의 얼음이 존재한다. 지지대의 역할을 하는 얼음의 경우, 빨간색으로 구분하여 볼 수 있으며 일반 얼음을 지탱해 주어 일반 얼음들이 깨지지 않도록 도와준다. 일반 얼음의 경우에는 1개의 지지대만이 연결되어 있어도 얼음이 깨지지 않지만 펭귄이 올라가 있는 얼음은 2개 이상의 지지대의 역할을 하는 얼음이 연결되어 있어야만 얼음이 깨지지 않는다. 이때, 지지대가 연결되어 있다는 것은 지지대로부터 서로 다른 일반 얼음들을 통해 연결 관계가 이어져 있는 것을 이야기한다. 특수 얼음 깨기 펭귄 게임에서 도도가 펭귄을 떨어뜨리지 않고 최대 몇 개의 얼음을 깰 수 있을까?

입력

첫째 줄에 얼음 블록의 개수 N(3≤N≤328000)과 지지대의 역할을 하게 되는 얼음의 개수 S(2≤S≤N−1), 펭귄이 위치한 얼음 블록의 번호 P(S<P≤N)가 주어진다. 지지대의 역할을 하게 되는 얼음의 개수가 S일 때, 1번부터 S번까지의 얼음은 지지대의 역할을 한다.

둘째 줄부터 N−1개의 줄에 두 개의 정수 A, B가 주어진다. 이는 A번 얼음과 B번 얼음이 연결되어 있음을 의미하며 같은 연결은 여러 번 주어지지 않는다.

게임 시작 시 펭귄은 일반 얼음 위에 위치해 있고 어떤 얼음도 깨지지 않은 상태로 시작하게 된다. 각 얼음들은 1번부터 N번까지 정수 번호로 주어져 있으며 서로 다른 두 얼음을 잇는 경로는 하나뿐이다. 더불어 서로 다른 지지대가 펭귄이 올라가 있는 얼음을 거치지 않고 연결되어 있는 경우는 없다.

 

출력

플레이어가 펭귄을 떨어트리지 않고 깰 수 있는 얼음의 최대 개수를 구하여라. 지지대의 역할을 하는 얼음 역시 깰 수 있는 얼음에 속한다.

---

펭귄이 있는 곳까지 거리를 재고 최소거리만 빼고 빼는 방식을 했는데 시간 초과...

import sys

# 총 얼음, 지지대, 펭귄얼음
N, S, P = map(int,sys.stdin.readline().split())
total = [[0] * (N+1) for _ in range(N+1)]

for _ in range(N-1):
    a, b = map(int,sys.stdin.readline().split())
    total[a][b] = 1
    total[b][a] = 1

lenP = [0] * (N+1)
visit = [0] * (N+1)
lenPs = []
for i in range(1,S+1):
    stack = [i]
    visit[i] = 1
    lenP[i] = 1
    while stack:
        a = stack.pop(0)
        for j in range(1,N+1):
            if total[a][j] and visit[j] == 0:
                if j == P:
                    lenPs.append(lenP[a])
                    stack = []
                    break
                visit[j] = 1
                lenP[j] = lenP[a] + 1
                stack.append(j)
lenPs.sort()
print(N - sum(lenPs[:2]) - 1)

인접 리스트를 사용하여 탐색하는 부분을 줄이니깐 빨라졌고 visit를 사용하여 길이를 구해 lenP를 삭제하였다.

import sys

# 총 얼음, 지지대, 펭귄얼음
N, S, P = map(int,sys.stdin.readline().split())
total = [[] for _ in range(N+1)]

# 인접 리스트
for _ in range(N-1):
    a, b = map(int,sys.stdin.readline().split())
    total[a].append(b)
    total[b].append(a)

visit = [0] * (N+1)
lenPs = []

# 각 정점에서 펭귄까지 거리 찾기
for i in range(1,S+1):
    # bfs 과정
    stack = [i]
    visit[i] = 1
    while stack:
        a = stack.pop(0)
        # **인접 리스트를 써서 이렇게하면 빠르게 실행가능**
        for j in total[a]:
            if visit[j] == 0:
                # 펭귄에 도착하면 거리를 저장한다.
                if j == P:
                    lenPs.append(visit[a])
                    # 펭귄에 도착하면 bfs 중단
                    stack = []
                    break
                visit[j] = visit[a] + 1
                stack.append(j)
# 펭귄까지 거리를 오름차순
lenPs.sort()
# 가장 짧은 거리의 얼음 2개와 펭귄의 얼음 빼고 다 뺄수있다.
print(N - sum(lenPs[:2]) - 1)