AtCoder が提供しているABC(AtCoder Beginner Contest)304 のE問題をC++とPythonで解いてみました。ABC304は、2023年6月3日21:00に実施されました。
この回は、サーバ攻撃を受けてジャッジ遅れが大きく発生したため、unrated となりました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Good Graph(Difficulty : 971)
問題はリンク先をご覧ください。
UnionFind木を使い、連結成分の頂点を記憶しておき、クエリを処理します。AtCoder Problems による Difficulty は 971 でした。
解答案
UnionFind とは
問題を解く前に UnionFind 木を紹介します。
以下のクラスは、「蟻本」(「プログラミングコンテストチャレンジブック」第2版 秋葉拓哉、岩田陽一、北川宣稔著、マイナビ 2012年)を参考に実装しました。ただし、後述する ACL (AtCoder Library) と合わせるため、メソッド名を変更しています。
class UnionFind {
private:
vector<int> par;
vector<int> rank;
public:
UnionFind(int n) {
par.resize(n);
rank.resize(n, 0);
for (int i = 0; i < n; i++) {
par[i] = i;
}
}
int leader(int x) {
if (par[x] == x) {
return x;
} else {
return par[x] = leader(par[x]);
}
}
void merge(int x, int y) {
x = leader(x);
y = leader(y);
if (x == y) {
return;
}
if (rank[x] < rank[y]) {
par[x] = y;
} else {
par[y] = x;
if (rank[x] == rank[y]) {
++rank[x];
}
}
}
bool same(int x, int y) {
return leader(x) == leader(y);
}
};
このクラスが提供しているメソッドは以下です。
メソッド名 | 説明 |
UnionFind(int n); | コンストラクタ、頂点の数 n を指定する。 |
int leader(int x); | 頂点 x の根を返す。 |
void merge(int x, int y) ; | 頂点 x を含む木と、頂点 y を含む木を併合する。 |
bool same(int x, int y); | 二つの頂点が同じ木に属している場合 true を、 それ以外の場合 false を返す。 |
UnionFind木は、初期状態では、N個の頂点に対して、N個の木があります。メソッド merge で2つの木を併合します。同じ木に属しているかをメソッド same で判断できます。それぞれの頂点に対して、メソッド leader は、所属している木の根を返します。
この問題はこのクラスを利用して解くことができます。
問題を解く方針
まず、N 個の頂点、M 個の辺で、UnionFind 木を構築します。
K 個の xi、yi は、この時点では、別々の根を持っているはずです。これは、「頂点 xi と yi を結ぶパスは存在しない」という制約があるためです。この根の pair を set コンテナに格納します。
Q 個の質問で与えられる pi と qi の根の pair が先ほどの格納した根の pair として存在していれば、この二つの頂点を結ぶと、xi と yi の木が連結となり、「良いグラフ」とはなりません。
C++ プログラム例(ABC304E)
上で述べた方針でプログラムにしました。
- 与えられた辺を merge する(57行目)。
- xi と yi の根の pair を set コンテナ leader に格納する。大小関係を考慮して、2通り格納しています(65、66行目)。
- 質問で与えられた pi と qi の根の pair が leader に存在していなければ、”Yes” を出力する。存在していれば “No” を出力する(74-77行目)。
以下が、最終的なC++プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
class UnionFind {
private:
vector<int> par;
vector<int> rank;
public:
UnionFind(int n) {
par.resize(n);
rank.resize(n, 0);
for (int i = 0; i < n; i++) {
par[i] = i;
}
}
int leader(int x) {
if (par[x] == x) {
return x;
} else {
return par[x] = leader(par[x]);
}
}
void merge(int x, int y) {
x = leader(x);
y = leader(y);
if (x == y) {
return;
}
if (rank[x] < rank[y]) {
par[x] = y;
} else {
par[y] = x;
if (rank[x] == rank[y]) {
++rank[x];
}
}
}
bool same(int x, int y) {
return leader(x) == leader(y);
}
};
int main()
{
int n, m;
cin >> n >> m;
UnionFind uf(n + 1);
for (int i = 0; i < m; ++i) {
int u, v;
cin >> u >> v;
uf.merge(u, v);
}
int k;
cin >> k;
set<pair<int, int>> leader;
for (int i = 0; i < k; ++i) {
int x, y;
cin >> x >> y;
leader.insert(make_pair(uf.leader(x), uf.leader(y)));
leader.insert(make_pair(uf.leader(y), uf.leader(x)));
}
int q;
cin >> q;
for (int i = 0; i < q; ++i) {
int x, y;
cin >> x >> y;
if (leader.find(make_pair(uf.leader(x), uf.leader(y))) == leader.end()) {
cout << "Yes" << endl;
} else {
cout << "No" << endl;
}
}
return 0;
}
AtCoder では、ACL (AtCoder Library) と呼ばれるライブラリを提供しています。ACL は多くの機能を含んでいます。UnionFind は、dsu (Disjoint Set Union) として提供されています。
以下は、ACL の dsu を使いました。変更したコードの背景色を変更しています。
#include <bits/stdc++.h>
#include <atcoder/dsu>
using namespace std;
using namespace atcoder;
int main()
{
int n, m;
cin >> n >> m;
dsu uf(n + 1);
for (int i = 0; i < m; ++i) {
int u, v;
cin >> u >> v;
uf.merge(u, v);
}
int k;
cin >> k;
set<pair<int, int>> leader;
for (int i = 0; i < k; ++i) {
int x, y;
cin >> x >> y;
leader.insert(make_pair(uf.leader(x), uf.leader(y)));
leader.insert(make_pair(uf.leader(y), uf.leader(x)));
}
int q;
cin >> q;
for (int i = 0; i < q; ++i) {
int x, y;
cin >> x >> y;
if (leader.find(make_pair(uf.leader(x), uf.leader(y))) == leader.end()) {
cout << "Yes" << endl;
} else {
cout << "No" << endl;
}
}
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC304E)
Python 版では、C++ の UnionFind を移植しました。UnionFind は、多くの方が実装しています。実装例のひとつとして読んでいただけたらと思います。
"""AtCoder Beginner Contest 304 E"""
class UnionFind():
def __init__(self, n):
self.n = n
self.par = []
self.rank = [0] * n
for i in range(self.n):
self.par.append(i)
def leader(self, x):
if self.par[x] == x:
return x
else:
self.par[x] = self.leader(self.par[x])
return self.par[x]
def merge(self, x, y):
x = self.leader(x)
y = self.leader(y)
if x == y:
return
if self.par[x] < self.par[y]:
self.par[x] = y
else:
self.par[y] = x
if self.rank[x] == self.rank[y]:
self.rank[x] += 1
def same(self, x, y):
return self.leader(x) == self.leader(y)
n, m = map(int, input().split())
uf = UnionFind(n + 1)
for i in range(m):
u, v = map(int, input().split())
uf.merge(u, v)
k = int(input())
leader = set()
for i in range(k):
x, y = map(int, input().split())
leader.add((uf.leader(x), uf.leader(y)))
leader.add((uf.leader(y), uf.leader(x)))
q = int(input())
for i in range(q):
x, y = map(int, input().split())
print("Yes" if not (uf.leader(x), uf.leader(y)) in leader else "No")
こちらも「AC」と判定されました。
最後に
この問題は、わたしの実力では難しかったですが、コンテスト中に解くことができました。このような経験があると、新しいデータ構造やアルゴリズムを学ぶ動機付けになります。
引き続き ABC の問題を紹介していきます。