AtCoder AGC 023 F번 01 on tree

작성일 | In

문제 링크

  • https://atcoder.jp/contests/agc023/tasks/agc023_f

문제 출처

  • AtCoder AGC 023 F번

사용 알고리즘

  • 그리디
  • 유니온 파인드

시간복잡도

  • $O(N log N)$

풀이

트리의 각 정점 안에 $a_i$개의 0과 $b_i$개의 1이 있다는 정보를 저장해줍시다.

$N$개의 정점 중 $\frac {a_i}{b_i}$가 큰 정점을 $v$, $v$의 부모를 $p$라고 하면 $p$ 바로 뒤에 $v$가 오는 것이 최적이라는 것을 exchange argument를 이용해 증명할 수 있습니다. $b_i$가 0인 경우에는 $\frac {a_i}{b_i}$를 무한대로 정의해주면 됩니다.
이 성질을 이용해 그리디한 알고리즘을 사용할 수 있습니다.

$\frac {a_i}{b_i}$가 가장 큰 정점인 $v$를 그의 부모 $p$ 바로 다음에 나오는 것이 최적이라는 것을 증명했기 때문에, 트리에서 $v$와 $p$를 merge해줄 수 있습니다!
노드를 merge해줄 때 merge하면서 발생하는 inversion을 모두 계산해준다면, merge한 이후에는 같은 정점 안에 있는 숫자에 대해서는 inversion을 신경쓰지 않고 노드와 노드 사이의 inversion만 신경을 써주면 됩니다.

두 노드 $p$와 $v$를 merge해줄 때 발생하는 inversion은 $b_p * a_v$입니다.
노드를 $N-1$번 merge해주면서 inversion을 계산해주면 정답을 구할 수 있습니다.

노드를 merge한 정보를 union-find로 관리해주고, $\frac {a_i}{b_i}$가 가장 큰 정점을 고르는 것을 priority queue등의 자료구조로 관리해주면 $O(N log N)$에 문제를 풀 수 있습니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;

struct Node{
	ll a, b, i;
	bool operator < (const Node &t) const {
		return a * t.b < t.a * b; // a/b < t.a/t.b
	}
}node[202020];

int par[202020];
void init(){
	for(int i=0; i<202020; i++) par[i] = i;
}
int find(int v){
	return v == par[v] ? v : par[v] = find(par[v]);
}

priority_queue<Node> pq;
int chk[202020], p[202020];
int n;

int main(){
	ios_base::sync_with_stdio(0); cin.tie(0);
	init();
	cin >> n;
	for(int i=2; i<=n; i++) cin >> p[i];
	for(int i=1; i<=n; i++){
		int t; cin >> t;
		if(t == 0) node[i] = {1, 0, i};
		else node[i] = {0, 1, i};
		if(i != 1) pq.push(node[i]);
	}
	ll ans = 0;
	while(pq.size()){
		auto t = pq.top(); pq.pop();
		int now = t.i;
		if(chk[now]) continue;
		chk[now] = 1;
		int nxt = find(p[now]);
		ans += node[now].a * node[nxt].b;
		node[nxt].a += node[now].a;
		node[nxt].b += node[now].b;
		if(nxt != 1) pq.push(node[nxt]);
		par[now] = nxt;
	}
	cout << ans;
}