【C++】等式方程式の充足可能性問題をUnion-Findで解く方法

問題概要

変数間の関係を表す方程式の配列が与えられます。各文字列 equations[i] は長さ4で、「a==b」または「a!=b」という2つの形式のいずれかを取ります。ここで a と b は小文字のアルファベット1文字からなる変数名を表します。求めるのは、与えられたすべての方程式を満たすように各変数に整数を割り当てることが可能かどうかです。可能な場合にのみ true を返します。

例えば、入力が ["a==b","b==c","a==c"] の場合、3つの変数すべてに同じ値を割り当てればよいため、答えは true になります。

アプローチ：Union-Find（素集合データ構造）

この問題は、Union-Find（Disjoint Set Union）と呼ばれるデータ構造を用いることで効率的に解けます。等号「==」で結ばれた変数同士は必ず同じグループに属し、一方で不等号「!=」で結ばれた変数同士は異なるグループに属していなければならないためです。

解法の手順

• まず、getParent() というメソッドを定義します。これは文字 x とマップ m を引数に取り、x が属するグループの親（代表元）を返します。
  • m[x] == x の場合は、そのまま x を返します。
  • そうでなければ、m[x] = getParent(m[x], m) として再帰的に親を辿り（経路圧縮）、m[x] を返します。
• メインの処理では、以下を実行します。
  • equal と notEqual の2つの配列を用意し、parent というマップを作成します。
  • n を e のサイズとします。
  • i を 0 から n-1 までループします。
    • parent[e[i][0]] = e[i][0]、parent[e[i][3]] = e[i][3] として各変数を初期化します。
    • e[i][1] が '=' ならインデックス i を equal 配列に、そうでなければ notEqual 配列に追加します。
  • i を 0 から equal.size()-1 までループし、等号の方程式を処理します。
    • index = equal[i] とし、u = e[index][0]、v = e[index][3] とします。
    • parent[getParent(u, parent)] = parent[getParent(v, parent)] として、両者の属するグループを統合します。
  • i を 0 から notEqual.size()-1 までループし、不等号の方程式を検証します。
    • index = notEqual[i] とし、u = e[index][0]、v = e[index][3] とします。
    • getParent(u, parent) == getParent(v, parent) の場合、同一グループ内で不一致が発生しているため false を返します。
  • すべてのチェックを通過できたら true を返します。

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

実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   char getParent(char x, map <char, char> m){
      if(m[x] == x) return x;
      return m[x] = getParent(m[x], m);
   }
   bool equationsPossible(vector<string>& e) {
      vector <int> equal;
      vector <int> notEqual;
      map <char, char> parent;
      int n = e.size();
      for(int i = 0; i < n; i++){
         parent[e[i][0]]= e[i][0];
         parent[e[i][3]]= e[i][3];
         if(e[i][1] == '='){
            equal.push_back(i);
         }else{
            notEqual.push_back(i);
         }  
      }
      for(int i = 0; i < equal.size(); i++){
         int idx = equal[i];
         char u = e[idx][0];
         char v = e[idx][3];
         parent[getParent(u, parent)] = parent[getParent(v, parent)];
      }
      for(int i = 0; i < notEqual.size(); i++){
         int idx = notEqual[i];
         char u = e[idx][0];
         char v = e[idx][3];
         if(getParent(u, parent) == getParent(v, parent)) return false;
      }
      return true;
   }
};
main(){
   vector<string> v1 = {"a==b","b==c","a==c"};
   Solution ob;
   cout << (ob.equationsPossible(v1));
}

入力

["a==b","b==c","a==c"]

出力

true

計算量について

方程式の数を N とすると、経路圧縮を施した getParent() の呼び出しはほぼ定数時間で完了するため、全体の時間計算量は O(N) 程度（std::map を使用する場合は O(N log N)）となります。また、変数は小文字アルファベット26種類に限られるため、空間計算量は O(1) とみなせます。