D - Pair of Balls Editorial by en_translator


Due to the Constraint that “there are exactly two balls of each color,” the following fact holds.

  • If there are multiple pairs of cylinders which have the same color of balls on the top, then we may perform the operation described in the problem statement in any order without changing the answer, that is, the result for the question “if it is possible to empty all the \(M\) cylinders.”

Therefore, in order to solve this problem, the following simulation is valid.

Repeat the following operation \(N\) times.

  • If no pair of cylinders has balls of the same color on the top, then output No and exit the program.

  • Otherwise, if there exists a pair of cylinders with balls of the same color, choose such pair arbitrarily and perform the operation described above.

If the program is not exit after the \(N\) repetitions, then output Yes.

The complexity of the simulation above is \(O(NM^2)\) or \(O(NM)\) naively. However, if you prepare an array \(\text{cnt}\) such that \(\text{cnt}[i]:=(\)the number of balls of color \(i\) that are on the top of any cylinders\()\), and while updating the array manage by a queue the set of integers \(i\) such that \(\text{cnt}[i]=2\), then the complexity can be improved to \(O(N+M)\).

Sample code (C++)

#include<bits/stdc++.h>
using namespace std;

int main(){
	int N,M; cin >> N >> M;
	queue<int> que;
	vector<queue<int>> A(M);
	vector<vector<int>> mem(N);
	for(int i=0; i<M; i++){
		int k; cin >> k;
		for(int j=0; j<k; j++){
			int a; cin >> a;
			A[i].push(a-1);
		}
		mem[A[i].front()].push_back(i);
	}
	for(int i=0; i<N; i++){
		if(mem[i].size() == 2){
			que.push(i);
		}
	}
	while(!que.empty()){
		int now = que.front(); que.pop();
		for(auto p: mem[now]){
			A[p].pop();
			if(!A[p].empty()){
				mem[A[p].front()].push_back(p);
				if(mem[A[p].front()].size() == 2){
					que.push(A[p].front());
				}
			}
		}
	}
	for(auto p: A){
		if(!p.empty()){
			cout << "No" << endl;
			return 0;
		}
	}
	cout << "Yes" << endl;
}

posted:
last update: