Official

E - Balls and Bag Query Editorial by en_translator


Let us consider the properties of each operation (especially, of first and second type).

Both operations increase or decrease the number of distinct numbers in the bag by one. Specifically, putting one ball into the bag increases the count by one or zero, and taking one ball from the bag decreases the count by one or zero.

Let us consider how we can exploit this property to process each operation fast. To determine if each operation changes the count or not, one need to find how many balls with a specific integer are in the bag. On insertion, if the bag already have a ball with that number, the number of distinct integers in the bag does not increase; otherwise it increases by one. Likewise, one can find the delta of the count on removal based on the number of balls with a specific integer written on them.

Thus, all that needed to process the queries is the current number of distinct integers and the number of balls with each integer in the bag. The former can be managed with an integer, and the latter with an array of length \(10^6\).

For more details, please refer to the sample code.

Sample code (C++):

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

int main() {
    int q;
    cin >> q;
    vector<int> cnt(1000000); // cnt[i]=(How many balls with (i+1) written on them are currently in the bag?)//
    int ans = 0; // The current number of distinct integers in the bag
    for (int i = 0; i < q; ++i) {
        int t;
        cin >> t;
        if (t == 1) {
            int x;
            cin >> x;
            x--;
            cnt[x]++;
            if (cnt[x] == 1) ans++; // In this case, there is only one ball with x written on it
        } else if (t == 2) {
            int x;
            cin >> x;
            x--;
            cnt[x]--;
            if (cnt[x] == 0) ans--; // In this case, there is no more ball with x written on it
        } else {
            cout << ans << "\n";
        }
    }
}

Sample code (Python):

q = int(input())
cnt = [0] * 1000000
ans = 0

for _ in range(q):
    t,*x = map(int,input().split())
    if t == 1:
        x[0] -= 1
        cnt[x[0]] += 1
        if cnt[x[0]] == 1:
            ans += 1
    elif t == 2:
        x[0] -= 1
        cnt[x[0]] -= 1
        if cnt[x[0]] == 0:
            ans -= 1
    else:
        print(ans)

posted:
last update: