🙆

ABC225 D - Play Train C++解答例

2021/10/31に公開約3,400字

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の形式で始まるクエリをクエリA2 x yの形式をクエリB3 xの形式をクエリCと呼ぶこととします。

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

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

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

クエリABの処理についての見通しが立ったので、次にクエリ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;
}

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

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

  2. https://atcoder.jp/contests/abc225/submissions/26928018 ↩︎

Discussion

ログインするとコメントできます