ABC226 C - Martial artist C++解答例

AtCoder Beginner Contest 226 C - Martial artistをC++で解きます。

問題

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

問題文

高橋くんは武術家です。武術家の覚えられる技は$N$個あり、技$1, 2, \ldots, N$と名前がついています。$1 \leq i \leq N$について、技$i$を習得するには時間$T_i$の修練が必要で、さらに、修練の開始時点で$A_{i, 1}, A_{i, 2}, \ldots, A_{i, K_i}$をすでに習得している必要があります。ここで、$1 \leq j \leq K_i$について、$A_{i, j} < i$であることが保証されます。

高橋くんは時刻0の時点で技を1つも覚えていません。高橋くんは同時に1つの技に対する修練しかできず、また、一度始めた修練を途中でやめることもできません。高橋くんが技$N$を習得するのに必要な時間の最小値を求めてください。

制約

• $1 \leq N \leq 2 \times 10^5$
• $1 \leq T_i \leq 10^9$
• $0 \leq K_i < i$
• $1 \leq A_{i, j} < i$
• $\sum_{i=1}^N K_i \leq 2 \times 10^5$
• $A_{i, 1}, A_{i, 2}, \ldots, A_{i, K_i}$はすべて異なる。
• 入寮は全て整数である。

解答例

依存関係はグラフで表現する

問題文からして複雑な問題ですが、とりあえず見通しを立てるところから始めます。

1から$N$までの番号がついた技があって、ある技$i(1\leq i \leq N)$を習得するためには$K_i$個の下位の技を覚えていなければならないという制約条件が課されています^[1]換えると、技$i$は0個以上の技$A_{i, 1}, A_{i, 2}, \ldots, A_{i, K_i}$に依存していると言えます。

このような頂点同士の依存関係を表現するには有向グラフを使うのが有効であると考えられます。

技$N$を覚えるために必要な全部の技は？

高橋くんの目的は最上位の技である技$N$を覚えることです。技$N$を覚えるために必要な技以外は覚える必要はありません。

技の依存関係をグラフで表したとき、$N$を覚えるために必要な技は、$N$から出ている有向辺が接続されている頂点集合であると言うことができます。また、そのそれぞれの頂点から伸びている有向辺が繋がっているノードもまた、$N$を覚えるために必要となります。

従って、頂点$N$から有向辺を辿って到達可能な全ての頂点を覚えなければ、技$N$は覚えられないということがわかります。

これらの技は頂点$N$からグラフを探索することで求めることができます。BFSでもDFSでもどちらでも可能です。

覚えるべき技を全部習得するのにかかる時間

次に、技$N$を覚えるために必要な技を全て習得するのにかかる時間を求めます。
便宜上、頂点$N$から到達可能な全てのノードの集合を$S$とし、$S$の各要素を番号が小さい順に$V_1, V_2, \ldots, V_{|S|}$と名付けることにします。

ここで、集合$S$に属する頂点$V_1, V_2, \ldots, V_{|S|}$のうち、最初に習得することができるのは番号が最も若い頂点$V_1$であることがわかります。
何故なら、技の依存関係は大きい番号から小さい番号にのみ伸びており、また集合$S$に属する頂点には必ず依存関係が存在するからです。

また、$V_1$を覚えた後に習得することができるのは、$V_1$の次に番号の若い頂点$V_2$です。
これを繰り返していくと、集合$S$に属する頂点(技)を全て覚えるためには、番号の若い順に習得していけばよいことがわかります。

実際に実装する

以上の考えを実際に実装します。問題を解くにあたって、2つのステップに分けて実装する必要があります。

1. $N$から到達可能な頂点を列挙する
2. 列挙した頂点を番号が小さい順に並べ、習得にかかる時間を足し合わせていく

1は簡単なグラフ探索で行えます。「$N$から到達可能か？」を表す配列を用意して探索を行うことで求めることができます。
1の結果を元に2の計算を行うことで、この問題を解くことができます。

実装例

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

c.cpp

#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
#include <queue>
 
int main() {
  int64_t N;
  std::cin >> N;
  std::vector<int64_t> T(N);
  std::vector<std::vector<int64_t>> graph(N);
  for (int64_t i = 0; i < N; i++) {
    int64_t t, k;
    std::cin >> t >> k;
    T.at(i) = t;
    for (int64_t j = 0; j < k; j++) {
      int64_t a;
      std::cin >> a;
      // 0-indexedでグラフを構築する
      graph.at(i).push_back(a - 1);
    }
  }
 
  // 頂点Nから到達可能な頂点を幅優先探索で探す
  std::vector<bool> reachable(N, false);
  // 頂点Nも含む
  reachable.at(N - 1) = true;
  std::queue<int64_t> candidate;
  candidate.push(N - 1);
 
  while (!candidate.empty()) {
    int64_t node = candidate.front();
    candidate.pop();
 
    for (int64_t next_node : graph.at(node)) {
      if (reachable.at(next_node)) continue;
      reachable.at(next_node) = true;
      candidate.push(next_node);
    }
  }
 
  // 番号が小さい順に習得所要時間を足し合わせていく
  int64_t answer = 0;
  for (int64_t i = 0; i < N; i++) {
    if (reachable.at(i)) {
      answer += T.at(i);
    }
  }
 
  std::cout << answer << std::endl;
 
  return 0;
}

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

おまけ：DFSによる解法

$N$から到達可能な頂点のうち、番号の若い順から習得所要時間を足し合わせるという計算を行えばこの問題を解くことができるため、構築したグラフに対して深さ優先探索を行うだけでもこの問題を解くことができます。

c.cpp

#include <iostream>
#include <vector>
#include <string>
#include <algorithm>
 
void dfs(std::vector<std::vector<int64_t>>& graph,
            std::vector<bool>& acquired,
            std::vector<int64_t>& T,
            int64_t& total_time,
            int64_t current_node)
{
  // 頂点iの技が依存している技を先に覚えに行く
  for (auto next_node : graph.at(current_node)) {
    if (acquired.at(next_node) == false) {
      dfs(graph, acquired, T, total_time, next_node);
    }
  }
  // 依存関係が解決したら技iを覚える
  acquired.at(current_node) = true;
  // 技iを覚えるのにかかった時間を足し合わせる
  total_time += T.at(current_node);
}
 
int main() {
  int64_t N;
  std::cin >> N;
  std::vector<int64_t> T(N);
  std::vector<std::vector<int64_t>> graph(N);
  for (int64_t i = 0; i < N; i++) {
    int64_t t, k;
    std::cin >> t >> k;
    T.at(i) = t;
    for (int64_t j = 0; j < k; j++) {
      int64_t a;
      std::cin >> a;
      graph.at(i).push_back(a - 1);
    }
  }
 
  // 探索済み(習得済み)かどうかを管理する配列
  std::vector<bool> acquired(N, false);
  int64_t answer = 0;
  dfs(graph, acquired, T, answer, N - 1);
 
  std::cout << answer << std::endl;
 
  return 0;
}

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

脚注

1. 4つ目の制約条件$1 \leq A_{i, j} < i$からこのことがわかります。ある技$i$を覚えようとしたとき、技$i$を覚えるために必要となる技のリスト$A_{i, 1}, A_{i, 2}, \ldots, A_{i, K_i}$の各要素$A_{i, j}(1 \leq j \leq K_i)$の値は$i$より小さくなるので、技$i$を覚えるために$i$より上位の技は関係ないということが示されています。