🙆

ABC225 D - Play Train C++解答例

2021/10/31に公開

C++

AtCoder Beginner Contest 225 D - Play TrainをC++で解きます。

問題

問題文をAtCoderのページより引用します。

問題文

高橋君は電車のおもちゃを連結させたり分離させたりして遊んでいます。
電車は $N$ 個あり、電車1、電車2、 $\ldots$ 、電車 $N$ と名前がついています。
はじめ電車どうしは連結しておらず全てバラバラです。

クエリが $Q$ 個与えられるので、与えられた順番に処理してください。
クエリは次の3種類のいずれかです。

1 x y：電車 $x$ の後部と、電車 $y$ の前部を連結させる。
以下のことが保証されます。

$x \neq y$

クエリを処理する直前に、電車 $x$ の後部と連結している電車は存在しない

クエリを処理する直前に、電車 $y$ の前部と連結している電車は存在しない

クエリを処理する直前に、電車 $x$ と電車 $y$ は異なる連結成分に属する

2 x y：電車 $x$ の後部と、電車 $y$ の前部を分離させる。
以下のことが保証されます。

$x \neq y$

クエリを処理する直前に、電車 $x$ の後部と電車 $y$ の前部は直接連結している。
3 x：電車 $x$ が含まれる連結成分に属する電車の番号を、先頭から順番に全て出力する。

制約

$2 \leq N \leq 10^5$

$1 \leq Q \leq 10^5$

$1 \leq x \leq N$

$1 \leq y \leq N$

入力はすべて整数

クエリは全て問題文の条件を満たす

3 xの形式のクエリで出力する電車の番号の個数の合計は $10^6$ 以下

解答例

各車両の接続情報を管理する

便宜上、1 x yの形式で始まるクエリをクエリ $A$ 、2 x yの形式をクエリ $B$ 、3 xの形式をクエリ $C$ と呼ぶこととします。

各車両について、前部についている車両の番号と後部についている車両の番号それぞれの情報を管理することを考えます。
これは、前部と後部でそれぞれ配列を用意するか、ペアを要素とする配列を1つ用意することで行えます。
車両が何も接続されていない場合は-1を記録しておくことで未接続を表現します。

この情報を配列で管理することで、クエリ $A$ 、 $B$ の両者を $\mathcal{O}(1)$ で処理することができます。

各連結成分の先頭の情報は持たなくてもよい

クエリ $A$ 、 $B$ の処理についての見通しが立ったので、次にクエリ $C$ について考えます。

Union-Find等を学んでいると、「その車両が属している連結成分の先頭車両の番号」も各車両に持たせたくなりますが、今回の問題ではそれをやる必要はありません。
先頭車両の情報を保持しないようにしたとき、連結成分に属する車両番号を先頭から列挙するためには、各車両を遡って先頭車両の番号を求めなければならなくなります。
一見するとこの操作に結構な計算量がかかりそうに見えますが、制約を見ると、「3 xの形式のクエリで出力する電車の番号の個数の合計は $10^6$ 以下」と書かれています^[1]。この制約条件より、クエリ $C$ の計算量は大きくなりすぎないことがわかるので、クエリ $C$ は愚直に実装しても問題ありません。

むしろUnion-Findライクに実装してしまうと、クエリ $A$ とクエリ $B$ の処理に余計な計算量がかかってしまい、TLEを引き起こしてしまいます^[2]。

実装例

C++実装例を以下に示します。

d.cpp

#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
 
int main() {
  int64_t N, Q;
  std::cin >> N >> Q;
  // 各車両の接続情報を管理する配列(1-indexed)
  std::vector<std::pair<int64_t, int64_t>> connection(N + 1, std::make_pair(-1, -1));
  for (int64_t i = 0; i < Q; i++) {
    int64_t t;
    std::cin >> t;
    if (t == 1) {
      // 車両xと車両yを連結させる
      int64_t x = 0;
      int64_t y = 0;
      std::cin >> x >> y;
      // 車両xの後部には車両yが付く
      connection.at(x).second = y;
      // 車両yの前部には車両xが付く
      connection.at(y).first = x;
    } else if (t == 2) {
      // 車両xと車両yを分離させる
      int64_t x = 0;
      int64_t y = 0;
      std::cin >> x >> y;
      // 分離した後は車両xの後部、および車両yの前部には何も接続しなくなる
      connection.at(x).second = -1;
      connection.at(y).first = -1;
    } else if (t == 3) {
      int64_t x = 0;
      std::cin >> x;
 
      // 前処理として、連結成分の先頭まで遡る
      while (connection.at(x).first != -1) {
        x = connection.at(x).first;
      }
      // 答えを格納する配列
      std::vector<int64_t> answer;
      // 先頭から末尾まで辿っていき、現れた番号を記録していく
      while (x != -1) {
        answer.push_back(x);
        x = connection.at(x).second;
      }
      // 答えを出力する
      std::cout << answer.size() << " ";
      for (size_t i = 0; i < answer.size(); i++) {
        if (i) std::cout << " ";
        std::cout << answer.at(i);
      }
      std::cout << std::endl;
    }
  }
 
  return 0;
}

実際に提出したコードはこちら。

脚注

この条件は「各クエリ $C$ で出力する番号の個数はそれぞれ $10^6$ 以下」という意味ではなく、「1つのテストケースでクエリ $C$ が合計 $L$ 個与えられ、それぞれで出力する番号の個数が $c_l$ 個 $(1 \leq l \leq L)$ であるとしたとき、総和 $\sum_{i = 1}^{L} c_i$ は $10^6$ 以下である」という意味だと思われます。 ↩︎
https://atcoder.jp/contests/abc225/submissions/26928018 ↩︎

問題

問題文

制約

解答例

各車両の接続情報を管理する

各連結成分の先頭の情報は持たなくてもよい

実装例

Discussion