【AtCoder】ABC 395 D - Pigeon Swap | 緑コーダーが解くAtCoder

atcoder.jp

配点: 350 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 919

問題概要
制約
考察
実装例
コメント

問題概要

それぞれ $1, 2, \cdots, N$ の番号が付いた $N$ 羽の鳩と巣があり、はじめ、鳩 $i$ は巣 $i$ にいる。

これらの鳩に対して、以下の3種類の形式のクエリを順に $Q$ 個処理せよ。

1 a b : 整数 $a,b \: (1\leq a, b\leq N)$ に対して、鳩 $a$ を今いる巣から取り出し、巣 $b$ へ移動する。
2 a b : 整数 $a,b \: (1\leq a, b\leq N)$ に対して、巣 $a$ にいる鳩をすべて巣 $b$ へ移動し、巣 $b$ にいる鳩をすべて巣 $a$ へ移動する。これら2つの移動は一斉に行われる。
3 a : 整数 $a \: (1\leq a\leq N)$ に対して、鳩 $a$ が今いる巣の番号を出力する。

制約

$1\leq N\leq10^6$
$1\leq Q\leq3\times10^5$
入力はすべて整数。

考察

クエリ処理の問題は、各クエリごとに必要な情報を書きだしていくとよい。なお、説明の都合上、クエリ1 → 3 → 2 の順に書いていく。

クエリ1

鳩 $i$ が今いる巣の番号が分かれば、それを $b$ に書き換えるだけ。したがって、

$\mathrm{pigeons}_i :=$ 鳩 $i$ が今いる巣の番号

という配列を用意する。

クエリ3

これも先ほどの $\mathrm{pigeons}_i$ という配列があれば、すぐに答えを出すことができる。

クエリ2

素直にやろうとすると

$\mathrm{pigeons}_i = a$ であるものを全て $\mathrm{pigeons}_i = b$ に書き換え、 $\mathrm{pigeons}_j = b$ であるものを全て $\mathrm{pigeons}_j = a$ に書き換える

となるが、 $\mathrm{pigeons}$ だけでこのような $i, j$ を探してくるためには、それぞれ $O(N)$ かかってしまう。

当然これでは TLE してしまうので、ここの処理を高速化したい。

ここで1つ思うのは、「鳩がいる巣を変えるのではなくて、巣の番号だけ変えてしまえばいいのでは？」ということ。

巣自体にも番号が書かれているので、巣に貼り紙(ラベル)を貼っておくことにして、そのラベルを貼り替えることを考える。

ということで、次の2つの配列を用意する。

$\mathrm{houses}_i :=$ 巣 $i$ に貼ってあるラベルの番号
$\mathrm{labels}_i :=$ ラベル $i$ が貼ってある巣の番号

「巣 → ラベル」と「ラベル → 巣」の関係をどちらも保持しておくことによって、クエリ1とクエリ3の両方に対応することができる。

各クエリの処理は以下の通り。なお、初期状態では $\mathrm{pigeons}_i = \mathrm{houses}_i = \mathrm{labels}_i = i \: (i = 1, 2, \cdots, N)$ である。

クエリ1 : $\mathrm{pigeons}_a = \mathrm{labels}_b$ とする。
クエリ2 : $\mathrm{labels}_a$ と $\mathrm{labels}_b$ をswapする。その後、 $\mathrm{houses}_{\mathrm{labels}_a} = a, \: \mathrm{houses}_{\mathrm{labels}_b} = b$ とする。
クエリ3 : $\mathrm{houses}_{\mathrm{pigeons}_a}$ を出力する。

実装例

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

#define all(x) (x).begin(), (x).end()

// ======================================== //

int main()
{
    int N, Q;
    cin >> N >> Q;

    vector<int> pigeons(N + 1), houses(N + 1), labels(N + 1);
    iota(all(pigeons), 0);
    iota(all(houses), 0);
    iota(all(labels), 0);

    while (Q--)
    {
        int op;
        cin >> op;

        if (op == 1)
        {
            int a, b;
            cin >> a >> b;

            pigeons[a] = labels[b];
        }
        else if (op == 2)
        {
            int a, b;
            cin >> a >> b;

            swap(labels[a], labels[b]);
            houses[labels[a]] = a;
            houses[labels[b]] = b;
        }
        else if (op == 3)
        {
            int a;
            cin >> a;

            cout << houses[pigeons[a]] << endl;
        }
    }

    return 0;
}