[BOJ] 1289번: 트리의 가중치

www.acmicpc.net/problem/1289

D[i] : i를 도착지로 하는 경로의 가중치의 합. 단, i의 subtree에서 i로 가는 경로에 대해서만 생각

으로 정의해보겠습니다. 그리고 root를 1로 두고, tree의 모양을 한번 생각해봅시다.

뭐 대충 이런식으로 1에 여러 개의 subtree가 달려있는 형태로 생겨있을 것입니다. 이 때 cost를 어떤식으로 구할 수 있을까요?

이 tree 내의 모든 경로는 1이 도착지 혹은 시작지인 경로 / 1이 도착지 혹은 시작지는 아니면서 1을 포함하는 경로 / 1을 포함하고 있지 않는 경로 로 나눌 수 있습니다.

i) 1이 도착지 혹은 시작지인 경로 : D[i]

ii) 1이 도착지 혹은 시작지는 아니면서 1을 포함하는 경로 : 1을 기준으로 서로 다른 subtree에 속해있는 두 node간의 경로. 이들의 합은 ((D[2]+1)*W[2])*((D[4]+1)*W[4])+((D[2]+1)*W[2])*((D[7]+1)*W[7])+((D[4]+1)*W[4])*((D[7]+1)*W[7]) 로 계산할 수 있습니다. 편의상 A[i] = (D[i]+1)*W[i] 로 쓴다면 A[2]*A[4]+A[4]*A[7]+A[2]*A[7]과 같습니다. 바로 와닿지 않는다면 간단한 예시에 대해 테스트해보면 쉽게 감이 올 것입니다.

iii) 1을 포함하지 않는 경로 : 1을 기준으로 같은 subtree에 속해있는 두 node간의 경로. 이들의 합은
D[2]+D[4]+D[7] 입니다.

이를 recursive하게 적용하면, D[1],D[2], ... , D[N]을 전부 cost에 더하고, 나의 자식 subtree에 대해 ii 번 과정과 같은 연산을 하면 됩니다. 그런데 ii) 번 과정에서 까딱하면 시간복잡도가 O(N^2)이 될 수 있습니다. subtree가 50000개라고 할 때, for문을 2번 돌며 합을 구한다고 치면 그렇게 됩니다. 그런데 간단한 트릭으로 시간복잡도를 O(N)으로 만들 수 있습니다.

바로 (ab+ac+ad+bc+bd+cd) = ((a+b+c+d)^2-(a^2+b^2+c^2+d^2))/2 라는 공식을 이용하는 것입니다.

중간중간 MOD를 취해주는 것과, int 범위를 뛰어넘을 수 있다는 것에 유의하면 답을 구해낼 수 있습니다.

 

github.com/encrypted-def/BOJ/blob/master/1289.cpp