村人の友好関係 (paizaランク S 相当) を C++17 で解いてみた (グラフ付き)

Posted at 2024-08-25

はじめに

たまに競技プログラミングを Rust で解いてみたという記事を書いている人です。

最近コードモンスターを進めています。コードモンスターの言語には Rust はなく、代わりにほかの言語を少しずつ触っています。

その中で ~~ガチャのためにダイヤが欲しく~~ Paiza スキルチェック問題を解いてみたところ、楽しめる感じでした。 ¹

ちょうどそんなときに、paiza×Qiita記事投稿キャンペーン「プログラミング問題をやってみて書いたコードを投稿しよう！」企画が行われていました。そこで便乗して 1問解いてみようと思いました。

本記事の概要

キャンペーンページの一番上にあった S ランク問題「村人の友好関係」を解きます
C++17 で実装します
- paizaプログラミングスキルチェックの各言語のバージョン、環境情報 | ITエンジニア向け転職・就活・学習サービス【paiza】
- C++20, C++23 は使えないようです ²
- Rust 1.59 も使えますが、コードモンスターがきっかけですので今回の記事では使いません
方法1: まず辺の数をそのままに multiset で実装します。1問制限時間をオーバーします
方法2: 辺の数を減らし、すべてのテストを制限時間内に通します
方法3: 辺の数を減らす方法を Disjoint-Set-Union (UnionFind) に切り替えて通します
感想

問題のイメージ

例題2 をグラフで描きます。

入力例:

出力例:

友好度が次のように決まっています。

+ 1: 最初に 1 さんが同好会に入りました。 🚩 は同好会メンバーを示します。同好会の内外を結ぶ線、言い換えると 🚩のある頂点とない頂点を結ぶ辺の最大値は 6 です。 6 を出力します。

- 1: 1 さんが同好会から抜けました。同好会の内外を結ぶ線の最大値は、線が存在しないため 0 とします。 0 を出力します。

+ 1: 再度 1 さんが同好会に入りました。同好会の内外を結ぶ線の最大値は 6 です。 6 を出力します。

+ 5: 5 さんが同好会に入りました。同好会の内外を結ぶ線の最大値は 9 です。 9 を出力します。

方針1: multiset でそのままシミュレーションする (1問制限時間超え)

このままシミュレーションします。内外を結ぶ辺の重みを状態 multiset で管理し、頂点 🚩 を追加削除するたびに辺も更新します。頂点追加のたびに multiset の最大値を調べれば、それが答えです。

操作	状態	最大値	追加	削除
	{}	0
+ 1			{6, 6}	{}
	{6, 6}	6
- 1			{}	{6, 6}
	{}	0
+ 1			{6, 6}	{}
	{6, 6}	6
+ 5			{9}	{6}
	{6, 9}	9

+ 同好会グループに参加するとき
- 辺でつながっている人がすでに同好会メンバー?
  - yes: 「片方だけ同好会」の条件から外れたので、辺の重みを「削除」する
  - no: 「片方だけ同好会」の条件になったので、辺の重みを「追加」する
- 同好会グループから抜けるするとき
- 辺でつながっている人がすでに同好会メンバー?
  - yes: 「片方だけ同好会」の条件になったので、辺の重みを「追加」する
  - no: 「片方だけ同好会」の条件から外れたので、辺の重みを「削除」する

このように実装します。

#include <iostream>
#include <set>
#include <vector>

using namespace std;

int main(void){
    int n, m, q;
    cin >> n >> m >> q;

    // 0-indexed な双方向グラフを組み立てる
    vector<vector<pair<int, int>>> graph(n); // [[(f, a), ...], ...]
    for (int i = 0; i < m; ++i) {
        int a, b, f;
        cin >> a >> b >> f;
        --a;
        --b;
        graph[a].push_back(make_pair(f, b));
        graph[b].push_back(make_pair(f, a));
    }

    // 同好会グループ nodes と、グループ内外を結ぶ edges を都度更新し、最大値を出力
    vector<bool> nodes(n);
    multiset<int> edges; // [f1, f2, ...]
    for (int i = 0; i < q; ++i) {
        char op;
        int q;
        cin >> op >> q;
        --q;
        if (op == '+') {
            nodes[q] = true;
            for (auto [f, p] : graph[q]) {
                if (nodes[p]) {
                    edges.erase(edges.find(f));
                } else {
                    edges.insert(f);
                }
            }
        } else if (op == '-') {
            nodes[q] = false;
            for (auto [f, p] : graph[q]) {
                if (nodes[p]) {
                    edges.insert(f);
                } else {
                    edges.erase(edges.find(f));
                }
            }
        }

        int f_max = 0;
        if (!edges.empty()) {
            f_max = *edges.rbegin();
        }
        cout << f_max << "\n";
    }

    return 0;
}

これで計算量は間に合うでしょうか。

2 ≦ N ≦ 500
1 ≦ M ≦ N*(N-1)/2
1 ≦ Q ≦ 50,000

制約から、最悪ケースでも頂点数 500, すべての辺数 125000, 1回の操作で辺の増減 500, 問い合わせ回数 50000 で $O(50000 \times 500 \times \log(125000))$ 。危ないけれどギリギリ通るんじゃないかなと実行してみたところ、

1問時間切れでした。 S問題は甘くなかった。

方針2: 辺の数を減らす

先ほどのように 1回の問い合わせに対して辺の追加削除を行っていると間に合いません。では、この問題で隣り合う辺すべてを見る必要があるでしょうか。

試しに、4人ですべての友好度が決まっている場合を考えてみます。友好度は適当に 2人の番号のかけ算としました。

4人が同好会に参加しているかどうか $2^4$ の組み合わせを調べます。

同好会メンバーかどうかを旗の色 🚩🏳️ で区別します。

人3 と人4 の旗の色が異なる場合、友好度の最大値は 12 です。これより大きな友好度の辺はありません。人 1, 人2 が同好会に参加していてもいなくても、結果は変わりません。

人3 と人4 の旗の色が同じ、かつ、人2 と人4 の旗の色が異なる場合、友好度の最大値は 8 です。旗の色が異なる頂点をつなぐ、8 より大きな友好度を持つ辺はありません。人 1 が同好会に参加していても、参加していなくても、結果は変わりません。

人2, 3, 4 の旗の色が同じ、かつ、人1 と人4 の旗の色が異なる場合、友好度の最大値は 4 です。旗の色が異なる頂点をつなぐ、4 より大きな友好度を持つ辺はありません。

人1, 2, 3, 4 の旗の色が同じ場合、友好度の最大値は 0 です。

これをまとめると、$2^4$ すべての組み合わせに対して、友好度の最大値になりうる辺は 3本しかないことが分かります。

これは、グラフからループを除いた木を作ることに対応します。重い辺から調べていき、辿ることができないものを追加していく感じです。

グループ	辺の重み	辺の頂点番号	異なるグループ？
{{1}, {2}, {3}, {4}}
	12	3, 4	✅
{{1}, {2}, {3, 4}}
	8	2, 4	✅
{{1}, {2, 3, 4}}
	6	2, 3	❌
	4	1, 4	✅
{{1, 2, 3, 4}}
	3	1, 3	❌
	2	1, 2	❌

詳しくは Minimum Spanning Tree (最小全域木) で検索を。辺の重みが負で、接続がないように見えるところは重み 0 の辺があるものとみなしたときの木に対応します。

グループを混ぜる処理は、頂点数が 500個以内と少ないこともあり、全頂点のグループ番号をループで調べて対象のものを置き換えるようにしました。

この構造ができると、全体に対する友好度の計算は $O(N)$ で行えます。たとえば同好会に 1, 2, 4, 3 の順に参加するとすると、

重み	頂点
12	(3, 4)
8	(2, 4)
4	(1, 4)

参加	同好会内	同好会外	対象の辺	友好度最大値
+1	1	2, 3, 4	(1, 4)	4
+2	1, 2	3, 4	(1, 4), (2, 4)	8
+4	1, 2, 4	3	(3, 4)	12
+3	1, 2, 3, 4			0

調べる組み合わせは少ないです。毎回全探索しても十分です。

という感じで書き直してみたものがこちら:

#include <algorithm>
#include <iostream>
#include <numeric>
#include <tuple>
#include <vector>

using namespace std;

int main(void){
    int n, m, q;
    cin >> n >> m >> q;

    // 重み -f の辺に対して、最小全域木 mst を組み立てる
    vector<tuple<int, int, int>> mst_edges; // [(f, a, b), ...]
    {
        vector<tuple<int, int, int>> all_edges; // [(f, a, b), ...]
        for (int i = 0; i < m; ++i) {
            int a, b, f;
            cin >> a >> b >> f;
            --a;
            --b;
            all_edges.push_back(make_tuple(f, a, b));
        }
        sort(all_edges.begin(), all_edges.end());
        reverse(all_edges.begin(), all_edges.end());

        vector<int> group(n); // 接続関係のあるグループ番号
        for (int i = 0; i < n; ++i) {
            group[i] = i;
        }

        for (auto [f, a, b] : all_edges) {
            int ga = group[a];
            int gb = group[b];
            if (ga == gb) {
                continue;
            }
            // b の属していたグループを a の属していたグループに混ぜる
            for (int i = 0; i < n; ++i) {
                if (group[i] == gb) {
                    group[i] = ga;
                }
            }
            mst_edges.push_back(make_tuple(f, a, b));
        }
    }

    // 最小全域木 mst の中で、2頂点の色が異なる辺のコスト最大値を出力する
    vector<bool> nodes(n);
    for (int i = 0; i < q; ++i) {
        char op;
        int q;
        cin >> op >> q;
        --q;
        if (op == '+') {
            nodes[q] = true;
        } else if (op == '-') {
            nodes[q] = false;
        }

        int f_max = 0;
        for (auto [f, a, b] : mst_edges) {
            if (nodes[a] != nodes[b]) {
                f_max = max(f_max, f);
            }
        }
        cout << f_max << "\n";
    }

    return 0;
}

今度は二分探索しません。最悪ケース計算から log が消え、頂点数 500, すべての辺数 125000, 問い合わせ回数 50000 で $O(50000 \times 500)$ となりました。さすがに通るでしょう。

お疲れさまでした。

方針3: Disjoint-Set-Union (UnionFind) を使う

競技プログラミングをしていると良く見る Disjoint-Set-Union (DSU, UnionFind) という構造があります。 2頂点が同じグループに属する = 辺を介して辿れるかを、高速に判定できます。解説ページが山ほどあります。

DSU の実装方法は速度や提供機能によっていろいろありますが、この問題では単純なもので十分です。AtCoder Library を参考にインターフェイスを用意しました。

class Dsu 
{
public:
	Dsu() = delete;
	Dsu(size_t n) {
        parents_.assign(n, -1);
	}

	int merge(int a, int b) {
		a = leader(a);
		b = leader(b);
		if (a != b) {
			parents_[b] = a;
		}
        return a;
	}

	bool same(int a, int b) {
		return leader(a) == leader(b);
	}

	int leader(int i) {
		if (parents_[i] < 0) {
			return i;
		}
        parents_[i] = leader(parents_[i]);
        return parents_[i];
	}

private:
	vector<int> parents_;
};

そして、先ほどのコードを Dsu を使うように書き換えます。

        vector<int> group(n); // 接続関係のあるグループ番号
        for (int i = 0; i < n; ++i) {
            group[i] = i;
        }

        for (auto [f, a, b] : all_edges) {
            int ga = group[a];
            int gb = group[b];
            if (ga == gb) {
                continue;
            }
            // b の属していたグループを a の属していたグループに混ぜる
            for (int i = 0; i < n; ++i) {
                if (group[i] == gb) {
                    group[i] = ga;
                }
            }
            mst_edges.push_back(make_tuple(f, a, b));
        }

上のコードを、下と入れ替えます。

        Dsu dsu(n);
        for (auto [f, a, b] : all_edges) {
            if (dsu.same(a, b)) {
                continue;
            }
            dsu.merge(a, b);
            mst_edges.push_back(make_tuple(f, a, b));
        }

コード全体

#include <algorithm>
#include <iostream>
#include <numeric>
#include <tuple>
#include <vector>

using namespace std;

class Dsu 
{
public:
	Dsu() = delete;
	Dsu(size_t n) {
        parents_.assign(n, -1);
	}

	int merge(int a, int b) {
		a = leader(a);
		b = leader(b);
		if (a != b) {
			parents_[b] = a;
		}
        return a;
	}

	bool same(int a, int b) {
		return leader(a) == leader(b);
	}

	int leader(int i) {
		if (parents_[i] < 0) {
			return i;
		}
        parents_[i] = leader(parents_[i]);
        return parents_[i];
	}

private:
	vector<int> parents_;
};

int main(void){
    int n, m, q;
    cin >> n >> m >> q;

    // 重み -f の辺に対して、最小全域木 mst を組み立てる
    vector<tuple<int, int, int>> mst_edges; // [(f, a, b), ...]
    {
        vector<tuple<int, int, int>> all_edges; // [(f, a, b), ...]
        for (int i = 0; i < m; ++i) {
            int a, b, f;
            cin >> a >> b >> f;
            --a;
            --b;
            all_edges.push_back(make_tuple(f, a, b));
        }
        sort(all_edges.begin(), all_edges.end());
        reverse(all_edges.begin(), all_edges.end());

        Dsu dsu(n);
        for (auto [f, a, b] : all_edges) {
            if (dsu.same(a, b)) {
                continue;
            }
            dsu.merge(a, b);
            mst_edges.push_back(make_tuple(f, a, b));
        }
    }

    // 最小全域木 mst の中で、2頂点の色が異なる辺のコスト最大値を出力する
    vector<bool> nodes(n);
    for (int i = 0; i < q; ++i) {
        char op;
        int q;
        cin >> op >> q;
        --q;
        if (op == '+') {
            nodes[q] = true;
        } else if (op == '-') {
            nodes[q] = false;
        }

        int f_max = 0;
        for (auto [f, a, b] : mst_edges) {
            if (nodes[a] != nodes[b]) {
                f_max = max(f_max, f);
            }
        }
        cout << f_max << "\n";
    }

    return 0;
}

こちらでも通ります。お疲れさまでした。

まあ、DSU を使ったところで、前半部の木作成処理 $O(N^2) = O(500^2)$ がだいたい $O(N) = O(500)$ になったというもので、後半部の検査 $O(50000 \times 500)$ の方が支配的です。

競技プログラマーとしてはこちらの方が慣れた型に沿っているかもしれない、くらいの感じで。

感想

普段 AtCoder アルゴリズムコンテストのお世話になっています。こちらと比べて、 Paiza スキルチェックを少し行っての感想です。経験が浅く、間違っているかもしれません。

全問正解しなくても正のスコアが出る
- 競技プログラミングのアルゴリズム系は、1つでも TLE したら 0点、制約内では何が来ても大丈夫なようにアルゴリズムを見直しましょう、という世界だと思っています
- 遅い方法でもスコアが出る、というのは新鮮でした
Paiza スキルチェック本番では一発勝負、試行錯誤できない
- 上の裏返しですが、「たぶんこれで通るかな」と思って提出、ダメなら 5分などのペナルティーをもらって再提出、みたいな方法が使えません
- 1回提出したらそれまで。解く時間より見直しが大事です
言語によって制限時間が異なる。スクリプト言語が不利にならない
- paizaプログラミングスキルチェックの各言語のバージョン、環境情報 | ITエンジニア向け転職・就活・学習サービス【paiza】 3秒から16秒まで
- AtCoder では言語によらず制限時間が同じです。そのため高速な言語 C++, Rust が人気そうです。 Python も人気が高いです
- Paiza スキルチェックでは、言語ごとの速度差をある程度制限時間で吸収しています。得意な言語で受ければ良さそうです
標準入出力例あり
- paizaプログラミングスキルチェックの値取得・出力サンプルコード | ITエンジニア向け転職・就活・学習サービス【paiza】
- 問題を解く前にまず困るところでして、このように情報がまとまっているのは嬉しいです ³
競技プログラミング用のライブラリーが使えない
- たとえば DSU はよく使う仕組みですから、手元にデッキとして持っておきたいものです
- Paiza スキルチェックの禁止事項「他人のコードを利用しない」に照らし合わせると、自分で書かないといけなさそうです
  - コーディングスキルチェック受験での禁止行為 – paizaFAQ
- 言い換えると、S 問題もデッキ不要な問題設定であって、アルゴリズムを知っているかどうかよりも実装力志向なのかな、と感じました
他の方の解答・解説があまり見られない
- AtCoder では、コンテスト直後に X(Twitter) での感想戦が盛り上がる、公式と有志の解説記事を読み比べできる、誤答も含めてすべての提出コードを誰でも見られる環境があります。参加するだけで力が付きそうです
- Paiza スキルチェックはそういうのがありません。試験という性質上、どうにもならないことかと思います
- 今回の企画のように、使用終了した問題を取り上げるのは面白いと思いました

ここまでで終わります。企画頂いた方、読んでいただいた方、どうもありがとうございました。

ガチャの結果、画像のように最高レアキャラの Perl がたくさん手に入りました。私は 20年以上前に Perl を書いていたこともありました。もうぜんぜん触っていません。 ↩
C++20 を使えないのが残念というのもありつつ、その前に C17++ 表記を見て「う……」と思ってしまいました。C17 と C++17 はけっこう違います。 ↩
Rust (beta) の入出力例も今後欲しいです。 proconio 無しで標準入力を受けるのはけっこう手間でして。 ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up