Codeforces Round #616 - C. Prefix Enlightenment

https://codeforces.com/contest/1290/problem/C

Problem - C - Codeforces

문제 깔끔해서 좋다.

 

램프 $N$개가 있고($1 \leq N \leq 3 \cdot 10^5$), 그들의 초기 On-Off 상태(1/0) 배열 $S$가 입력으로 주어진다. 집합 $A_1, A_2, \dots , A_k \subseteq \{1, 2, 3, \dots, N\}$이 있다.($1 \leq K \leq 3 \cdot 10^5$) 단, 임의의 $1 \leq i < j < k \leq N$을 고르면 $A_i \cap A_j \cap A_k = \emptyset$을 만족한다. $A_i$ 집합을 선택하면 그 집합의 원소에 해당하는 모든 램프가 toggle 된다.

 

$N$ 줄에 거쳐서 $i$번째 줄에 1번부터 $i$번 램프까지 전부 켜지도록 선택해야할 부분집합 $A_x$ 들의 최소 개수를 출력하여라. 단, 모든 전구를 동시에 켜는 것이 가능한 경우만 입력으로 들어온다.

 

 

못 풀었다. B번 까지 40분 안에 풀어서 충분히 노려볼만했다고 생각했는데, 아직 내 실력에는 어려운 문제였던 것 같다. 그래도 덕분에 많이 배웠다.

 

에디토리얼도 좋았지만 투어리스트의 코드가 정말 멋있었다.

 

 

$a$번 램프를 켜고 끌 수 있는 부분집합 번호의 모음을 $B[a]$라고 하자. 문제의 조건에 의해 $|B(a)| \leq 2$ 이다.

 

쿼리의 순서에 맞게 램프의 번호를 $1$부터 $N$까지 루프를 돌자. (램프 번호 변수를 $a$라고 하자.)

• $|B(a)| = 0$ 인 경우, 애초에 할 수 있는 동작이 없다. 직전에 출력한 값과 똑같이 출력한다.
• $|B(a)| = 1$ 인 경우, $B(a) = \{x\}$라고 하자. $A_x$의 선택/미선택 여부가 하나로 고정된다.
• $|B(a)| = 2$ 인 경우, $B(a) = \{x, y\}$라고 하자. $x$와 $y$ 사이에 간선을 추가해보자. 그래프가 나온다.

만약 새로 연결된 간선을 추가했을 때, DFS를 통해 연결된 컴포넌트를 순회해서 선택/미선택 여부가 결정된 노드가 등장한다면, 그 컴포넌트에 있는 전체 노드의 선택/미선택 여부가 결정된다. 그런 노드가 등장하지 않는 상황이 문제다.

 

그런 노드들을 Union-Find 자료구조로 관리하자. 한 컴포넌트에 있는 노드들의 선택/미선택 여부가 고정이 되지는 않았지만, 노드들 간의 일치/불일치 여부는 결정되어 있다.

 

무슨 말이냐고 하면.. $|B(a)| = 2$, $B(a) = \{x, y\}$이고, 간선 $\{x, y\}$를 추가한다고 생각하자. 만약 $S[a] = 1$이라면 노드 $x$, $y$의 선택/미선택 여부는 일치한다. 왜냐하면 둘 다 동시에 선택되거나 동시에 선택되지 않아야지만 $a$번째 램프의 번호가 $1$로 유지되기 때문이다. 아니고 $S[a] = 0$이라면 노드 $x$, $y$의 선택/미선택 여부는 일치하지 않는다. 둘 중 하나만 선택되어야만 $a$번째 램프의 번호가 $1$이 되기 때문이다.

 

여기까지는 생각했었다. 구현에서 막혔다. 유니온-파인드 자료구조에서 루트가 아닌 두 노드를 합칠 때, 컴포넌트 노드 전체의 일치/불일치 여부를 어떻게 조합해야할 것이냐가 제일 문제였다. 여기다가 선택/미선택이 확실한 노드가 나올 때마다 DFS를 돌려주는 코드를 넣으니 스파게티가 돼서 고칠 엄두가 나지가 않았다.

 

 

그러나 대회가 끝난 직후 투어리스트의 코드를 보고 감탄했다. 애초에 DFS를 사용하지도 않았다.. ㄷㄷ

 

유니온-파인드 자료구조에서 컴포넌트의 각 노드들이 부모 노드와 선택/미선택 값이 일치/불일치하는지에 대한 배열 $diff\_with\_parent$을 저장하자.

 

유니온-파인드 자료구조에서 find 함수를 실행하면 재귀를 통해 루트까지 거슬러 올라간다. path compression을 통해 루트까지의 경로 상에 위치한 모든 노드들의 부모가 루트가 되도록 갱신해줄 수 있다. 그 과정에서 $diff\_with\_parent$도 xor을 이용해서 갱신해준다. 이렇게 되면 find 함수로 불러졌던 어떤 노드가 컴포넌트의 루트 노드와 선택/미선택 여부가 일치하는지 일치하지 않는지를 바로 알 수 있게 된다.

 

그렇게 되면 union 함수에서 두 노드를 합칠 때 $S[a]$ 비트 값과 $diff\_with\_parent$ 배열 값에 따라 다르게 합치면 된다.

 

또한, 루트 노드에서 컴포넌트의 노드들 중 루트와 선택/미선택 여부가 일치하는 것들의 개수와 일치하지 않는 것의 개수를 저장하고 union 할 때 이 값도 갱신해준다.

 

이 컴포넌트에서 골라야하는 부분집합의 최소 개수는 루트와 일치하는 노드의 개수, 루트와 일치하지 않는 노드의 개수 중 최소값에 해당한다.

 

Truth value가 고정된 노드와 병합된 경우 컴포넌트 전체의 truth value도 고정된다. 이 여부도 유니온-파인드 구조체에서 관리한다.

 

이렇게 되면 DFS를 하지 않아도 컴포넌트 전체 값을 고정시키는 동작을 $O(log^*n)$만에 수행 가능하다.

 

 

약간 세그먼트 트리의 lazy propagation같은 개념을 유니온 파인드의 path compression에 접목한 코딩 테크닉이라고 해야하나..

 

암튼 find 재귀함수로 부모노드와의 관계를 path compression을 통해 루트와의 관계로 연결시킬 수 있다는 점을 배운건 정말 큰 것 같다.

 

역시 코딩은 뚝배기가 깨져가면서 배우는게 제일 좋다!

 

 

코드

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef long double ld;

#define pb push_back
#define eb emplace_back

#define CPPIO ios::sync_with_stdio(0); cin.tie(0); cout << std::setprecision(10); cout << fixed

#define forn(x, n) for (int x=0; x<(n); x++)
#define for1(x, n) for (int x=1; x<=(n); x++)
#define ford(x, n) for (int x=(n)-1; x>=0; x--)
#define ford1(x, n) for (int x=(n); x>=1; x--)
#define forr(x, n, m) for (int x=(n); x<=(m); x++)
#define forrd(x, n, m) for (int x=(m); x>=(n); x--)

class union_find_size {
	public:
		int max_node;
		vector<int> parents;
		vector<int> eq, neq, diff, fixed, tfval;

		union_find_size(int max_node) {
			this->max_node = max_node;
			parents.assign(max_node+1, 0);
			eq.assign(max_node+1, 1);    // 루트와 TF값 같은 노드 개수
			neq.assign(max_node+1, 0);   // 루트와 TF값 다른 노드 개수
			diff.assign(max_node+1, 0);  // 부모와 TF값이 같은지 다른지
			fixed.assign(max_node+1, 0); // TF값이 고정되어있는지
			tfval.assign(max_node+1, 0); // 고정된 TF값 배열
			for (int i=0; i<=max_node; i++)
				parents[i] = i;
		}

		int find(int a) {
			if (parents[a] == a) return a;
			int fa = find(parents[a]);
			diff[a] = diff[parents[a]] ^ diff[a]; // from tourist's code
			fixed[a] = fixed[parents[a]];
			if (fixed[a]) {
				tfval[a] = tfval[parents[a]] ^ diff[a];
			}
			parents[a] = fa; // path compression
			return fa;
		}

		void fix(int a, int tf) {
			int fa = find(a);
			tfval[fa] = tf ^ diff[a];
			fixed[fa] = 1;
		}

		bool unio(int a, int b, int df) {
			find(a);
			find(b);
			df = df ^ diff[a] ^ diff[b];
			if ((a=find(a)) != (b=find(b))) {
				if (df) {
					eq[b] += neq[a];
					neq[b] += eq[a];
				} else {
					eq[b] += eq[a];
					neq[b] += neq[a];
				}
				diff[a] = df;
				parents[a] = b;
				if (fixed[a]) fix(a, tfval[a]);
				fixed[b] = fixed[a] || fixed[b];
				return 1;
			} else return 0;
		}

		int value(int a) {
			int fa = find(a);
			if (fixed[fa]) {
				if (tfval[fa]) return eq[fa];
				else return neq[fa];
			}
			return min(eq[fa], neq[fa]);
		}
};

const int MAX_N = 3e5 + 20;

int N, K;
string S;
union_find_size Uf(MAX_N);
int Cnt;
vector<int> B[MAX_N];


int main() {
	CPPIO;
	
	cin >> N >> K;
	cin >> S;
	S = "_" + S;

	for1(i, K) {
		int q;
		cin >> q;
		while (q--) {
			int a;
			cin >> a;
			B[a].pb(i);
		}
	}

	for1(a, N) {
		int lamp = S[a] - '0';
		int sz = B[a].size();
		if (sz == 0) {
		} else if (sz == 1) {
			int x = B[a][0];
			Cnt -= Uf.value(x);
			Uf.fix(x, !lamp);
			Cnt += Uf.value(x);
		} else {
			int x = B[a][0], y = B[a][1];
			if (Uf.find(x) != Uf.find(y)) {
				Cnt -= Uf.value(x);
				Cnt -= Uf.value(y);
				Uf.unio(x, y, !lamp);
				Cnt += Uf.value(x);
			}
		}
		cout << Cnt << endl;
	}

	return 0;
}