けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder AGC 002 B - Box and Ball (400 点)

うまくシミュレーションする系

問題へのリンク

問題概要

N 個の箱があります。 箱は 1 から N まで番号が振られています。 最初、1 番目の箱には赤いボールが 1 個入っています。 また、2~N 番目の箱には白いボールが 1 個ずつ入っています。

高橋君は M 回の操作を順に行います。 i 回目の操作では、xi 番目の箱から適当なボールを 1 個選び、それを yi 番目の箱へ移します。

高橋君がすべての操作を終えた後、赤いボールが入っている可能性のある箱は何個か求めてください。

制約

  •  2 \le N \le 10^{5}
  •  1 \le M \le 10^{5}

考えたこと

本当に上手にシミュレーションすれば通る。まず、どのように赤いボールを動かしたとしても、各ターンにおける各箱のボールの個数は同一なことに注目する。それに注目して

  • nums[ v ] := 各ターン終了時の箱 v のボールの個数
  • can[ v ] := 各ターン終了時に箱 v に赤いボールが入っている可能性があるかどうか

としてあげて、これを毎ターン更新する。箱 x から箱 y へボールを移すときの nums の更新は自明。単純に nums[x]--, nums[y]++ で OK。can の更新は基本的には

  • can[x] = true だったら、can[y] = true に

とすればよいのだけど、注意点として、

  • x から y へ移したときに x が空になるなら can[x] = false に

というのが必要。

#include <iostream>
#include <vector>
using namespace std;

int main() {
    int N, M;
    cin >> N >> M;
    vector<int> nums(N, 1), can(N, 0);
    can[0] = 1;
    for (int i = 0; i < M; ++i) {
        int a, b; cin >> a >> b; --a, --b;
        if (can[a]) can[b] = 1;
        if (nums[a] == 1) can[a] = 0;
        nums[a]--;
        nums[b]++;
    }
    int res = 0;
    for (int i = 0; i < N; ++i) if (can[i] && nums[i]) ++res;
    cout << res << endl;
}