[BOJ_Platinum 4] 24520 - Meet in the middle [python]
24520번: Meet In The Middle
첫 번째 줄에 마을의 수 $N$, 약속의 수 $K$가 주어진다. $(1 \le N, K \le 100\,000)$ 이어지는 줄부터 $N-1$개의 줄에 도로 정보를 나타내는 세 정수 $u$, $v$, $w$가 주어진다. $u$번 마을과 $v$번 마을 사이에
www.acmicpc.net
문제
신촌 왕국에는 번부터 번까지 개의 마을이 있고 개의 도로가 서로 다른 두 마을을 연결하고 있다. 모든 마을은 도로를 통해 다른 마을로 갈 수 있고, 이때 마을 내에서의 이동 거리는 무시할 수 있다.
현재 신촌 왕국은 "대칭병"의 대유행으로 혼란을 겪고 있다. 대칭병은 어떤 것이든 대칭을 이루지 않으면 심신이 불편해져 일상생활을 할 수 없게 만드는 병이다. 대칭병의 여파는 엄청났는데 심지어 사람과 사람 간의 만남조차 어렵게 됐다. 신촌 왕국의 두 사람이 만나는 상황을 생각해보자. 약속 장소가 둘 중 한 사람의 마을에 더 가깝다면 아주 불편해진다. 따라서 두 사람이 사는 각 마을까지의 거리가 정확히 같은 곳을 약속 장소로 정해야 할 것이다. 즉, Meet in the middle. 가운데에서 만나야 한다!
신촌 왕국 도로 정보와 만나려는 두 사람의 마을이 주어졌을 때, 약속 장소로 적절한 위치를 구해보자.
입력
첫 번째 줄에 마을의 수 , 약속의 수 가 주어진다.
이어지는 줄부터 개의 줄에 도로 정보를 나타내는 세 정수 , , 가 주어진다. 번 마을과 번 마을 사이에 길이 의 도로가 있다는 뜻이다., u≠v,
이어지는 줄부터 개의 줄에 만나려는 두 사람의 마을 번호 , 가 주어진다.
출력
개의 줄에 두 사람의 만날 수 있는 마을의 번호를 출력한다. 그러한 마을이 없다면 -1을 출력한다. 가능한 답이 여러 가지라면, 두 사람이 사는 각 마을까지 거리의 합이 짧은 것을 출력한다.
==================================
2022 ICPC Sinchon Winter Algorithm Camp Contest - 중급 난이도의 E번 문제로 출제된 문제입니다.
LCA(최소공통조상)과 이분탐색, 희소배열을 통해서 풀 수 있습니다.
각 마을과 길로 이루어진 트리에 대해서, 임의의 노드를 루트 노드로 설정해 깊이 우선탐색을 진행하며 희소배열을 완성한 뒤, 입력으로 주어지는 쿼리를 O(logN)의 시간 복잡도로 처리하면 되는 문제입니다.
잘 생각해보면, 두 마을 사이의 최단거리는 LCA를 기준으로 A마을에서부터 LCA까지의 거리, B마을에서 LCA까지의 거리의 합임을 알 수 있습니다.
A마을에서 LCA까지의 거리를 $Dist_{A-LCA}$
B마을에서 LCA까지의 거리를 $Dist_{B-LCA}$라고 하면,
$Dist_{A-LCA} + Dist_{B-LCA}$ 가 홀수일 경우 가운데에서 만나는 경우는 존재할 수 없음을 알수있습니다.
두 사람이 가운데에서 만나기 위해선, 서로 동일한 거리를 이동하기 위해 총 거리가 짝수여야합니다.
하지만 약속장소가 A쪽으로 X만큼 가까워지면 B의 입장에서는 X만큼 멀어지게되고, 결국 총 거리는 동일하게 유지되기 때문에, 총 거리가 홀수인 경우 가운데에서 만날 수 없습니다.
또한, A마을에서 B마을까지 거리의 최솟값은 항상 $Dist_{A-LCA} + Dist_{B-LCA}$ 입니다.
위의 내용은 다음과 같이 귀류법으로 증명할 수 있습니다.
i) LCA가 아닌 더 높은 조상(CA)을 기준점으로 설정한 거리가 LCA기준 거리보다 더 짧다.
$Dist_{A-LCA} + Dist_{B-LCA} + 2Dist_{B-LCA} > Dist_{A-LCA} + Dist_{B-LCA}$ 이므로 모순
ii) LCA가 아닌 마을 A(혹은 B)의 자식을(C) 기준점으로 설정한 거리가 LCA기준 거리보다 더 짧다.
$Dist_{A-LCA} + Dist_{B-LCA} + 2Dist_{A-C} > Dist_{A-LCA} + Dist_{B-LCA}$ 이므로 모순
따라서 언제나 LCA를 기준으로 거리를 책정하는 것이 가장 짧은 거리임을 알 수 있습니다.
아래 소스와 같이 희소배열을 구현하고, 두 마을의 LCA, 그리고 LCA까지의 거리를 구합니다.
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**5)
#log16까지만 필요함 2^16+a
n,tc = map(int,input().split())
graph = [[] for _ in range(n+1)]
depth = [1 for i in range(n+1)]
sparseParent = [[0 for _ in range(17)] for _ in range(n+1)] #노드번호,distance
sparseDist = [[0 for _ in range(17)] for _ in range(n+1)] #노드번호,distance
depth[0] = 0
#최대 (3log10만) * 10민
def makeTree(node,parent,cost):
if not parent:
depth[node] = 0
else:
depth[node] = depth[parent] + 1
sparseParent[node][0] = parent
sparseDist[node][0] = cost
for i in range(1,17): #log로 형성하기
sparseParent[node][i] = sparseParent[sparseParent[node][i-1]][i-1] #희소 부모 배열
sparseDist[node][i] = sparseDist[node][i-1] + sparseDist[sparseParent[node][i-1]][i-1] #희소거리 배열
if not sparseParent[node][i]: break
for nextNode,cost in graph[node]:
if nextNode == parent: continue
makeTree(nextNode,node,cost)
def getLca(a,b):
diff = depth[a] - depth[b]
lcaDist_a = 0
lcaDist_b = 0
if diff:
for i in range(17):
if not (diff >> i) & 1: continue
lcaDist_a += sparseDist[a][i] #거리 증가시키기
a = sparseParent[a][i] #부모로 끌어올리기
if depth[a] == depth[b]: break
if a == b: return [a,lcaDist_a,lcaDist_b]
#올렸을 때 동일하다면 리턴
curA,curB = a,b
for i in range(16,-1,-1):
if not sparseParent[curA][i] or sparseParent[curA][i] == sparseParent[curB][i]: continue
lcaDist_a,lcaDist_b = lcaDist_a + sparseDist[curA][i],lcaDist_b + sparseDist[curB][i]
curA,curB = sparseParent[curA][i],sparseParent[curB][i]
return [sparseParent[curA][0],lcaDist_a + sparseDist[curA][0],lcaDist_b + sparseDist[curB][0]]
A에서 약속장소까지의 거리 = $Dist_A$
B에서 약속장소까지의 거리 = $Dist_B$ = B에서 lca까지의 거리 + lca에서 약속장소까지의 거리 이므로,
희소배열을 사용해 둘의 거리가 같아지는 시점을 O(logN)만에 구할 수 있습니다. ( 이진탐색의 원리와 동일합니다.)
def findEqulDist(a,b):
if depth[a] < depth[b]: a,b = b,a
lca,distA,distB = getLca(a,b) #이렇게 두개 받아오기
#print(lca, distA, distB)
if distA == distB: return lca
if (distA + distB) % 2 : return -1
if distA < distB: distB,distA,a,b = distA,distB,b,a
cur = a
aDist = 0
for i in range(16,-1,-1):
if not sparseParent[cur][i]: continue #범위 밖인 경우 패스
nextADist = aDist + sparseDist[cur][i]
if nextADist > distB + (distA - nextADist): continue
#if nextADist == distB + (distA - nextADist): return sparseParent[cur][i]
aDist = nextADist
cur = sparseParent[cur][i]
return cur if aDist == distB + (distA - aDist) else -1
전체 소스는 아래와 같습니다.
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**5)
#log16까지만 필요함 2^16+a
n,tc = map(int,input().split())
graph = [[] for _ in range(n+1)]
depth = [1 for i in range(n+1)]
sparseParent = [[0 for _ in range(17)] for _ in range(n+1)] #노드번호,distance
sparseDist = [[0 for _ in range(17)] for _ in range(n+1)] #노드번호,distance
depth[0] = 0
def makeTree(node,parent,cost): #희소배열 구현
if not parent:
depth[node] = 0
else:
depth[node] = depth[parent] + 1
sparseParent[node][0] = parent
sparseDist[node][0] = cost
for i in range(1,17): #log로 형성하기
sparseParent[node][i] = sparseParent[sparseParent[node][i-1]][i-1] #희소 부모 배열
sparseDist[node][i] = sparseDist[node][i-1] + sparseDist[sparseParent[node][i-1]][i-1] #희소거리 배열
if not sparseParent[node][i]: break
for nextNode,cost in graph[node]:
if nextNode == parent: continue
makeTree(nextNode,node,cost)
def getLca(a,b): #두 마을 사이의 lca구하기
diff = depth[a] - depth[b]
lcaDist_a = 0
lcaDist_b = 0
if diff: #둘의 높이를 같게 하기
for i in range(17) :
if not (diff >> i) & 1: continue
lcaDist_a += sparseDist[a][i] #거리 증가시키기
a = sparseParent[a][i] #부모로 끌어올리기
if depth[a] == depth[b]: break
if a == b: return [a,lcaDist_a,lcaDist_b]
curA,curB = a,b
for i in range(16,-1,-1):#Binary Lift로 LCA구하기
if not sparseParent[curA][i] or sparseParent[curA][i] == sparseParent[curB][i]: continue
lcaDist_a,lcaDist_b = lcaDist_a + sparseDist[curA][i],lcaDist_b + sparseDist[curB][i]
curA,curB = sparseParent[curA][i],sparseParent[curB][i]
return [sparseParent[curA][0],lcaDist_a + sparseDist[curA][0],lcaDist_b + sparseDist[curB][0]]
def findEqulDist(a,b): #약속 장소 구하기
if depth[a] < depth[b]: a,b = b,a
lca,distA,distB = getLca(a,b) #이렇게 두개 받아오기
#print(lca, distA, distB)
if distA == distB: return lca
if (distA + distB) % 2 : return -1
if distA < distB: distB,distA,a,b = distA,distB,b,a
cur = a
aDist = 0
for i in range(16,-1,-1):
if not sparseParent[cur][i]: continue #범위 밖인 경우 패스
nextADist = aDist + sparseDist[cur][i]
if nextADist > distB + (distA - nextADist): continue
aDist = nextADist
cur = sparseParent[cur][i]
return cur if aDist == distB + (distA - aDist) else -1
for _ in range(n-1):
f,t,c = map(int,input().split())
graph[f].append([t,c])
graph[t].append([f,c])
makeTree(1,0,0)
for _ in range(tc):
a,b = map(int,input().split())
print(findEqulDist(a,b))