【AtCoder】ABC 391 C - Pigeonhole Query | 緑コーダーが解くAtCoder

atcoder.jp

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

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

問題概要

$1$ から $N$ までの番号が付いた $N$ 匹の鳩と、 $1$ から $N$ までの番号が付いた $N$ 個の巣がある。はじめ、鳩 $i$ は巣 $i$ にいる。

$Q$ 個のクエリが与えられるので順に処理せよ。クエリは2種類あり、以下のいずれかの形式で与えられる。

1 P H : 鳩 $P$ を巣 $H$ に移動させる。
2 : 複数の鳩がいる巣の個数を出力する。

制約

$2\leq N\leq 10^6$
$1\leq Q\leq 3\times 10^5$
$1\leq P,H\leq N$
1つ目のクエリについて、鳩 $P$ は移動する前に巣 $H$ にいない。
入力は全て整数。

考察

とりあえず、現在鳩がどの巣にいるのかを管理したいので、そのための配列pigeonsを用意する。

また、巣に何匹の鳩がいるのかも知りたいので、そのための配列homesを用意する。

$\mathrm{pigeons}_i :=$ 鳩 $i$ が現在いる巣の番号
$\mathrm{homes}_i :=$ 巣 $i$ に現在いる鳩の数

こうすることで、クエリ1は以下のように書ける。

homes[pigeons[P]]--; // 鳩 P が今いる巣を離れる

homes[H]++; // 鳩 P が巣 H に移動

pigeons[P] = H; // 鳩の位置も更新

クエリ2では $\mathrm{homes}_i \leq 2$ となる $i \: (1 \leq i \leq N)$ の個数を数えたいが、このために毎回forループを回していると TLE してしまう。

そこで、この個数自体を予め変数に持っておくことにする(ここではcntとした)。

cntが更新されるのは以下のタイミングとなる。

$\mathrm{homes}_{\mathrm{pigeons}_{P}} = 1$ となったとき、 $\mathrm{cnt} \leftarrow \mathrm{cnt} - 1$
$\mathrm{homes}_{H} = 2$ となったとき、 $\mathrm{cnt} \leftarrow \mathrm{cnt} + 1$

あとは、クエリ2ではcntを出力すればよいだけ。

結局、どちらのクエリも $O(1)$ で処理できることになった。

実装例では 0-indexed になっていることに注意。

実装例

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

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

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

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

    int cnt = 0;
    vector<int> pigeons(N), homes(N, 1);
    iota(all(pigeons), 0);
    while (Q--)
    {
        int t;
        cin >> t;

        if (t == 1)
        {
            int P, H;
            cin >> P >> H;
            P--;
            H--;

            homes[pigeons[P]]--;
            if (homes[pigeons[P]] == 1)
                cnt--;

            homes[H]++;
            if (homes[H] == 2)
                cnt++;

            pigeons[P] = H;
        }

        else if (t == 2)
        {
            cout << cnt << endl;
        }
    }

    return 0;
}