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

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

AtCoder ABC 298 C - Cards Query Problem (茶色, 300 点)

クエリ形式の問題に慣れたり、クエリに答えるためのデータ構造を考えたりすることに慣れたりするための練習問題!

問題概要

1 から  N までの番号のついた  N 個の箱と、何も書かれていない無数のカードがある。今、 Q 個のクエリが与えられるので、順に処理せよ。

  • クエリタイプ 1 (1 i j):カードを 1 枚選んで数  i を書き込み、そのカードを箱  j に入れる
  • クエリタイプ 2 (2 i):箱  i に入っているカードに書かれた数を昇順ですべて答えよ
    • ただし、同じ箱に同じ数のカードが複数枚ある場合は、その枚数分出力する
  • クエリタイプ 3 (3 i):数  i が書かれたカードが入っている箱の番号を昇順ですべて答えよ
    • ただし、同じ数のカードが同じ箱に複数枚入っている場合は、箱の番号を一度だけ出力する

制約

  •  1 \le N, Q \le 2 \times 10^{5}
  • カードの数は  1 以上  2 \times 10^{5} 以下

解法:クエリ処理問題とは

初めてクエリ処理問題を解こうとする人にとっては、今回の問題は戸惑ったかもしれない。クエリ処理問題とは

  • データを効率的に管理する「データ構造」を適切に設計した上で
  • データ構造」を更新したり (今回のクエリタイプ 1)
  • データ構造」に対する問いに答えたり (今回のクエリタイプ 2, 3)

するような問題のことをいう。

多くの場合、クエリは  10^{5} 回以上投げられるので、1 回 1 回のクエリには素早く答えることが必要となる。クエリに素早く答えられるようなデータ構造を適切に設計することが肝となる。

今回のデータ構造:2 次元配列

とはいえ、今回の問題では、データ構造に関する高度な知識は必要としない。比較的単純な「2 次元配列」で実現できる。

まず、クエリタイプ 3 を一旦無視して、クエリタイプ 1, 2 を処理する方法を考えよう。そのためには


どの箱に、どんな数字のカードが入っているか


という情報を効率的に管理できればよい。具体的には、下図のように 2 次元配列 box2card を用意して、box2card[i] が箱  i に入っているカードの番号の集合を表すようにすることで実現できる。

そして、クエリタイプ 1, 2 に対しては次のようにすればよい。

クエリタイプ 1 の処理

「数  i の書かれたカードを箱  j に入れる」という処理は、次のように書ける。

box2card[j].push_back(i)

(C++)

box2card[j].append(i)

(Python3)

クエリタイプ 2 の処理

「箱  i に入っているカードを昇順に出力する」という処理は次のように書ける。

sort(box2card[i].begin(), box2card[i].end());  // 並び替え
for (auto x : box2card[i]) cout << x << " ";
cout << endl;

(C++)

# 並び替えて、各要素を文字列にして、空白文字で join して出力
print(' '.join(map(str, sorted(box2card[i]))))

(Python3)

クエリタイプ 3 を考慮する

ここまでで、クエリタイプ 1, 2 を処理できたので、最後にクエリタイプ 3 を処理する。

クエリタイプ 3 では、「箱を指定してカードの数字を出力する」のではなく、「数字を指定して、その数字のカードの入った箱を出力する」ことが要求される。

そのため、2 次元配列 box2card と同様の考え方で、次のような配列を用意しよう。


  • card2boxcard2box[i] が「数  i を含む箱の集合」を表すような配列

そしてクエリタイプ 3 では「重複は除く」ということなので、単純な 2 次元配列 (各要素が配列であるような配列) よりも、「各要素が集合型であるような配列」の方が都合がいい。(ただし、単純な 2 次元配列でも実装可能)

C++ であれば、vector<set<int>> 型が使える。Python3 であれば、各要素が set 型であるような list が使える。

具体的な実装については、以下の解答例を参考に。

解答例

以上の解法は次のように実装できる。

C++

#include <bits/stdc++.h>
using namespace std;
const int MAX = 210000;  // 箱番号やカードの数がこれ以上にはならない

int main() {
    int N, Q;
    cin >> N >> Q;
    
    // box2card[i] := 箱 i に入っているカードたち
    vector<vector<int>> box2card(MAX);
    
    // card2box[i] := 数 i のカードの入った箱たち
    vector<set<int>> card2box(MAX);
    
    // 各クエリ処理
    for (int _ = 0; _ < Q; ++_) {
        int type, i, j;
        cin >> type;
        
        if (type == 1) {
            // クエリタイプ 1
            cin >> i >> j;
            
            // 数 i のカードを箱 j に追加
            box2card[j].push_back(i);
            card2box[i].insert(j);
        } else if (type == 2) {
            // クエリタイプ 2
            cin >> i;
            
            // 箱 i に入っているカードを並び替える
            sort(box2card[i].begin(), box2card[i].end());
            
            // 出力
            for (auto x : box2card[i]) cout << x << " ";
            cout << endl;
        } else {
            // クエリタイプ 3
            cin >> i;
            
            // 出力
            for (auto x : card2box[i]) cout << x << " ";
            cout << endl;
        }
    }
}
            

Python3

# 箱番号やカードの数がこれ以上にはならない
MAX = 210000

# 入力
N = int(input())
Q = int(input())

# box2card[i] := 箱 i に入っているカードたち
box2card = [[] for _ in range(MAX)]
    
# card2box[i] := 数 i のカードの入った箱たち
card2box = [set() for _ in range(MAX)]
    
# 各クエリ処理
for _ in range(Q):
    query = list(map(int, input().split()))
    if query[0] == 1:
        # 数 i のカードを箱 j に追加
        i, j = query[1], query[2]
        box2card[j].append(i);
        card2box[i].add(j);
    elif query[0] == 2:
        i = query[1]
        
        # 箱 i に入っているカードを並び替えて出力
        print(' '.join(map(str, sorted(box2card[i]))))
    else:
        i = query[1]
        
        # 数 j の書かれたカードの入っている箱を出力
        print(' '.join(map(str, sorted(card2box[i]))))