C++で解くコーススケジュールIV ― 前提条件クエリの判定方法

問題概要

受講できるコースが全部で n 個あり、各コースには 0 から n-1 までの番号が割り当てられています。

一部のコースには直接の前提条件が存在します。たとえば、コース 0 を受講する前にコース 1 を受講しておく必要がある場合、この関係はペア [1,0] として表現されます。

ここで、コース数 n、直接の前提条件ペアのリスト、そしてクエリペアのリストが与えられます。

各クエリ queries[i] に対して、「コース queries[i][0] はコース queries[i][1] の前提条件であるか」を判定してください。最終的に、すべてのクエリへの回答をブーリアン値のリストとして返します。

注意すべき点として、前提条件の関係は推移的です。つまり、コース a がコース b の前提条件であり、さらにコース b がコース c の前提条件であるならば、コース a はコース c にとっても前提条件となります。

たとえば、入力が n = 3、prerequisites = [[1,2],[1,0],[2,0]]、queries = [[1,0],[1,2]] のとき、出力は [true, true] となります。

解法の考え方

この問題はトポロジカルソート(幅優先探索)を用いて効率的に解けます。基本的なアイデアは次のとおりです。

• 各コースについて、「そのコースを受講するために必要なすべての先行コース」を集合として記録します。
• トポロジカル順にノードを処理しながら、隣接ノードへ祖先(先行コース)の情報を伝播させていきます。
• この前計算を行っておけば、各クエリには集合の検索だけで即座に回答できます。

アルゴリズムの手順

1. 定数 N := 110 を定義します。
2. 結果を格納する配列 ret を用意します。
3. 各ノードの入次数を管理するマップ in を定義します。
4. v の各要素 it について、次を行います。
   • graph[it[0]] の末尾に it[1] を追加します。
   • in[it[1]] を 1 増やします。
5. キュー q を定義します。
6. i := 0 から i < n の範囲で i を 1 ずつ増やしながら、in[i] が 0 のノード i を q に挿入します。
7. q が空になるまで、レベルごとに幅優先探索を繰り返します。
   • sz := q のサイズ とし、sz が 0 になるまで次を繰り返します。
     • node := q の先頭要素 を取り出し、キューから削除します。
     • graph[node] の各要素 it について:
       • in[it] を 1 減らします。
       • c[node] 内の全要素 x を c[it] に挿入します。
       • node 自身も c[it] に挿入します。
       • in[it] が 0 になったら、it を q に挿入します。
8. x の各要素 it について、c[it[1]] における it[0] の存在確認(count の結果)を ret の末尾に追加します。
9. ret を返します。

C++実装例

理解を深めるために、以下の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<bool> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]"<<endl;
}
const int N = 110;
class Solution {
   public:
   vector <int> graph[N];
   map <int, set <int>> c;
   vector<bool> checkIfPrerequisite(int n, vector<vector<int>>& v, vector<vector<int>>& x) {
      vector<bool> ret;
      map<int, int> in;
      for (auto& it : v) {
         graph[it[0]].push_back(it[1]);
         in[it[1]]++;
      }
      queue<int> q;
      for (int i = 0; i < n; i++) {
         if (in[i] == 0)
            q.push(i);
      }
      map<int, int> idx;
      for (int lvl = 1; !q.empty(); lvl++) {
         int sz = q.size();
         while (sz--) {
            int node = q.front();
            q.pop();
            for (auto& it : graph[node]) {
               in[it]--;
               for (auto& x : c[node])
                  c[it].insert(x);
               c[it].insert(node);
               if (in[it] == 0) {
                  q.push(it);
               }
            }
         }
      }
      for (auto& it : x) {
         ret.push_back(c[it[1]].count(it[0]));
      }
      return ret;
   }
};
main(){
   Solution ob;
   vector<vector<int>> prerequisites = {{1,2},{1,0},{2,0}}, queries = {{1,0},{1,2}};
   print_vector(ob.checkIfPrerequisite(3, prerequisites, queries));
}

入力

3, {{1,2},{1,0},{2,0}}, {{1,0},{1,2}}

出力

[1, 1]

まとめ

この手法では、各コースの祖先(前提条件)集合を事前に構築しておくため、クエリごとにグラフを再探索する必要がありません。計算量は集合のマージを含めておおよそ O(n³) オーダー、必要なメモリは O(n²) であり、n ≤ 100 程度の制約下では十分高速に動作します。前提条件が推移的である点を正しく扱えるのが、このアプローチの大きな強みです。