LeetCode 399 Evaluate Division

題意

給出多個形如 a/b=k 的式子以及詢問 a/b=?，要求對回答每個詢問的結果，若從式子推導不出詢問的答案，則返回-1。其中 a, b 爲任意字符串，k 爲實數。例如：

式子：a / b = 2.0, b / c = 3.0
詢問：a / c = ?, b / a = ?, a / e = ?, a / a = ?, x / x = ? 
答案：6.0, 0.5, -1.0, 1.0, -1.0 

思路

不難想到的是將除法變成乘法的形式，例如 a/b=2 可以變爲 a=2b，那麼通過將每個字符串表示爲其它字符串的倍數，就能得到任意兩個之間的倍數關係。例如對於 a/b=2, c/b=3, e/c=3，可以得到 a=2b, c=3b, e=9b 那麼 a/b, a/c, b/c, b/e, c/e 就都能得到。進一步，這種關係可以通過並查集這種數據結構實現，以類似於樹的關係將 e 表示成根節點 b 的倍數，即可判斷任意兩個之間的關係。與並查集類似，在建立關係樹的時候要考慮兩個子樹之間的合併關係：合併子樹的根。步驟如下：

1. 構造數據結構以表示a=kb這種關係
2. 考慮每一個等式 a/b=k
   1. a, b 都不在任一子樹中，以 b 爲根建立子樹
   2. a, b 其中一個在某個子樹中， 找到子樹的根，將不在子樹上的加入子樹
   3. a, b 在不同的子樹上，找到兩顆子樹的根，合併子樹
3. 對每個詢問，通過與根的倍數關係得出結果

代碼

struct relation {
    int who;
    double times;
    relation(int w, double t) {
        who = w, times = t;
    }
};
class Solution {
private:
    map<string, int> ids;
    vector<relation> fathers;
    int idtot = 0;
public:
    vector<double> calcEquation(vector<vector<string>> &equations, vector<double> &values, vector<vector<string>> &queries) {
        for (int i = 0; i < equations.size(); i++) {
            setFather(equations[i], values[i]);
        }
        vector<double> resvec;
        for (auto &&vec : queries) {
            double ret = queryK(vec[0], vec[1]);
            resvec.push_back(ret);
        }
        return resvec;
    }
    relation findFaa(int id) {
        relation ret = fathers[id];
        if (ret.who == id)
            return ret;
        relation tmp = findFaa(ret.who);
        ret.who = tmp.who;
        ret.times *= tmp.times;
        return ret;
    }
    void setFather(vector<string> equation, double value) {
        string a = equation[0], b = equation[1];
        if (ids.find(a) == ids.end() && ids.find(b) == ids.end()) {
            int aid = idtot++;
            int bid = idtot++;
            ids[a] = aid;
            ids[b] = bid;
            fathers.push_back(relation(bid, value));
            fathers.push_back(relation(bid, 1));
        } else if (ids.find(a) == ids.end()) {
            int aid = idtot++;
            int bid = ids[b];
            ids[a] = aid;
            relation bfa = findFaa(bid);
            fathers.push_back(relation(bfa.who, value * bfa.times));
        } else if (ids.find(b) == ids.end()) {
            int aid = ids[a];
            int bid = idtot++;
            ids[b] = bid;
            relation afa = findFaa(aid);
            fathers.push_back(relation(afa.who, (1.0 / value) * afa.times));
        } else {
            int aid = ids[a], bid = ids[b];
            relation afa = findFaa(aid), bfa = findFaa(bid);
            if (afa.who != bfa.who) {
                fathers[afa.who].times = bfa.times * value / afa.times;
                fathers[afa.who].who = bfa.who;
            }
        }
    }
    double queryK(string a, string b) {
        if (ids.find(a) == ids.end() || ids.find(b) == ids.end())
            return -1;
        int aid = ids[a], bid = ids[b];
        relation afa = findFaa(aid), bfa = findFaa(bid);
        if (afa.who != bfa.who)
            return -1;
        return afa.times / bfa.times;
    }
};