[17969] Gene Tree

https://www.acmicpc.net/problem/17969

17969번: Gene Tree

 

Tree DP로 풀 수 있는 문제입니다. 편의상 자식노드에 포함된 모든 리프 노드들과 현재 부모노드와의 거리 제곱의 합을 $\sum_{i=1}^{n} a_i^2$, 거리의 합을 $\sum_{i=1}^{n} a_i$, 리프노드의 개수를 $n$이라고 놓겠습니다.

 

자식노드에서 부모노드로 리프노드들을 가져오는 과정에서 거리가 $k$씩 추가됩니다. 따라서 거리 제곱의 합과 거리의 합을 각각 $\sum_{i=1}^{n} (a_i+k)^2$, $\sum_{i=1}^{n} (a_i+k)$로 갱신시켜줘야 합니다.

먼저 거리 제곱의 합의 식을 풀어서 써보면 다음과 같습니다.

 

$\sum_{i=1}^{n} (a_i+k)^2 = \sum_{i=1}^{n}  a_i^2 + 2k \sum_{i=1}^{n} a_i + \sum_{i=1}^{n}k^2$

 

$\sum_{i=1}^{n} a_i^2$과 $\sum_{i=1}^{n} a_i$, 그리고 $n$의 값 모두 알고 있기 때문에 식의 값을 $O(1)$에 갱신시켜주는 것이 가능합니다. 거리의 합 역시

 

$\sum_{i=1}^{n} (a_i+k) = \sum_{i=1}^{n} a_i + \sum_{i=1}^{n} k$

 

이렇게 식을 풀어볼 수가 있고 거리제곱과 마찬가지로 $O(1)$에 갱신 가능합니다. 

 

이제 서로 다른 자식노드들의 리프노드 거리의 제곱 합을 구해서 답에 더해줘야합니다. 두 자식노드들 간의 거리 제곱의 합은 

 

$\sum_{i=1}^{n} \sum_{j=1}^{m} (a_i+b_j)^2$

 

가 되며, 이 식을 풀어서 써본다면

 

$\sum_{i=1}^{n} \sum_{j=1}^{m} (a_i+b_j)^2 = \sum_{i=1}^{n} \sum_{j=1}^{m} a_i^2 + 2\sum_{i=1}^{n} \sum_{j=1}^{m}a_ib_j + \sum_{i=1}^{n} \sum_{j=1}^{m}b_j^2 \\ = m\sum_{i=1}^{n}a_i^2 + 2\sum_{i=1}^{n}a_i \sum_{j=1}^{m}b_j + n\sum_{j=1}^{m}b_j^2$

 

가 됩니다. 위의 식의 값을 구하기 위해 필요한 값들을 이미 알고 있으므로 이 역시 $O(1)$ 시간에 계산이 가능합니다.

매번 서로 다른 자식 노드 2개를 골라 저 식을 계산하면 $O(n^2)$의 시간이 걸리지만, 누적합을 사용한다면 $O(n)$에 가능해하며, 충분히 시간내로 가능합니다. 

이제 부모노드한테 리프노드들을 건네주면 됩니다. 위에서 말한 내용들을 각 노드들에서 실행하면 문제에서 요구하는 답을 얻어낼 수 있습니다.