C++でUnion-Findを使って遺伝子の突然変異グループ総数を求める方法

問題概要

同じ長さの文字列リスト genes があるとします。各要素は A、C、G、T のいずれかの文字のみで構成されており、次のルールが成り立ちます。

• 2つの文字列 s1 と s2 がちょうど1文字だけ異なる場合、s1 と s2 は同じ突然変異グループに属します。
• s1 と s2 が同じグループに属し、s2 と s3 が同じグループに属するならば、s1 と s3 も同じグループに属します(推移性)。

このとき、生成できる突然変異グループの総数を求めるのが目的です。

具体例

たとえば、入力が genes = ["ACGT", "ACGC", "ACTT", "TTTT", "TGTT"] の場合、出力は 2 になります。これは、次の2つの突然変異グループが存在するためです。

• ["ACGT", "ACGC", "ACTT"]
• ["TTTT", "TGTT"]

解法のアプローチ:Union-Find(素集合データ構造)

この問題は、Union-Find(Disjoint Set Union)を用いることで効率的に解けます。ポイントは、すべてのペアを比較する代わりに、「各文字列から1文字だけ変えた候補」を列挙し、それがリスト内に存在するかを高速に判定する点です。

1. parent マップの定義: 各文字列の親ノードを記録するマップ parent を用意します。
2. getPar() 関数: 経路圧縮を行いながら、ある要素の根(グループの代表元)を再帰的に求めます。
3. unite() 関数: 2つの要素が属する集合を統合します。統合が行われた場合は true、すでに同じ集合だった場合は false を返します。
4. ok() 関数: 2つの文字列を比較し、異なる文字がちょうど1個であるかどうかを判定します。
5. メイン処理(solve):
   • まず配列 v をソートし、set を作成して O(log N) での存在判定を可能にします。
   • ret の初期値は v のサイズ(初期状態ではすべての文字列が独立したグループ)です。
   • 各文字列の各位置 j に対して、その文字を A / C / G / T の残り3文字に置き換えた文字列 temp を生成し、temp が set 内に存在すれば unite() で統合します。統合に成功するたびに ret を1減らします。

最終的な ret が、求める突然変異グループの総数になります。

C++での実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    map<string, string> parent;

    string getPar(string& a){
        if(parent[a] == a)
            return a;
        return parent[a] = getPar(parent[a]);
    }

    bool unite(string& a, string& b){
        string parA = getPar(a);
        string parB = getPar(b);
        if(parA != parB){
            parent[parA] = parB;
            return true;
        }
        return false;
    }

    bool ok(string &a, string& b){
        int cnt = 0;
        for(int i = 0; i < a.size(); i++){
            cnt += a[i] != b[i];
        }
        return cnt == 1;
    }

    int solve(vector<string> v) {
        sort(v.begin(), v.end());
        set<string> s(v.begin(), v.end());

        int ret = v.size();
        for(auto& it : v){
            parent[it] = it;
        }
        for(auto& it : v){
            for(int j = 0; j < it.size(); j++){
                string temp = it;
                for(char x : {'A', 'C', 'G', 'T'}){
                    if(x != it[j]){
                        temp[j] = x;
                        if(s.count(temp)){
                            if(unite(temp, it)) ret--;
                        }
                    }
                }
            }
        }
        return ret;
    }
};

int main(){
    vector<string> v = {"ACGT", "ACGC", "ACTT", "TTTT", "TGTT"};
    Solution ob;
    cout << ob.solve(v);
}

入力

{"ACGT", "ACGC", "ACTT", "TTTT", "TGTT"}

出力

2

計算量の目安

N を文字列の個数、L を文字列の長さとすると、各文字列につき L × 3 個の候補を生成し、それぞれについて set への存在判定(O(log N))と Union-Find 操作(ほぼ定数時間)を行います。したがって全体の時間計算量は O(N × L × log N)、空間計算量は O(N × L) となります。全ペアを直接比較する O(N² × L) の手法に比べ、N が大きい場合に大きく高速化できる点がこの手法の利点です。