Official

C - Pigeonhole Query Editorial by en_translator


This problem can be solved by dynamically managing

  • which nest each pigeon is,
  • how many pigeons are in each nest, and
  • how many nests with multiple pigeons are there.

Let \(\mathrm{pos}[i]\) be the nest where pigeon \(i\) is, \(\mathrm{cnt}[i]\) be the number of pigeons in nest \(i\), and \(\mathrm{ans}\) be the number of nests with multiple pigeons. Query \(2\) asks \(\mathrm{ans}\).

Consider query \(1\). After pigeon \(P\) moves, there are two cases where \(\mathrm{ans}\) may be affected:

  • the number of pigeons in its old nest changes from \(2\) to \(1\);
  • the number of pigeons in its new nest changes form \(1\) to \(2\).

One can update \(\mathrm{pos}, \mathrm{cnt}, \mathrm{ans}\) according to the move of pigeon \(P\), considering these cases. For more details, please refer to the sample code.

N, Q = map(int, input().split())
ans = 0
cnt = [0] + [1] * N
pos = list(range(N + 1))
for i in range(Q):
    query = list(map(int, input().split()))
    if query[0] == 1:
        P, H = query[1:]
        # Pigeon P is originally in nest pos[P].
        # If the number of pigeons in nest post[P] decreases from 2 to 1 by moving, decrease ans by one
        if cnt[pos[P]] == 2:
            ans -= 1
        # Decrement the number of pigeons in nest pos[P] by one
        cnt[pos[P]] -= 1

        # Change pigeon P's nest to H
        pos[P] = H
        # Increment the number of pigeons in nest H by one
        cnt[pos[P]] += 1
        # If the number of pigeons in nest H increases from 1 to 2 by moving, increase ans by one
        if cnt[pos[P]] == 2:
            ans += 1
    else:
        print(ans)

posted:
last update: