除法求值 带权并查集 力扣399

除法求值题目

class Solution {
public:
    int find(int p, vector<int> &father, vector<double> & weight) 
    {
        if (father[p] != p) {
            int tf = father[p]; //保存原父节点
            father[p] = find(father[p], father, weight);
            weight[p] = weight[p] * weight[tf]; //更新权值
        }
        return father[p];
    }

    void merge(int a, int b, double w, vector<int> &father, vector<double> & weight)
    {
        int fa = find(a, father, weight);
        int fb = find(b, father, weight);
        if (fa == fb)
            return;

        father[fa] = b;
        weight[fa] = w / weight[a];
    }

    vector<double> calcEquation(vector<vector<string>>& equations, vector<double>& values, vector<vector<string>>& queries) 
    {
        unordered_map<string, int> points;
        int cnt = 0;

        for (auto item : equations) {
            if (points.find(item[0]) == points.end()) {
                points[item[0]] = cnt++;
            }
            if (points.find(item[1]) == points.end()) {
                points[item[1]] = cnt++;
            }
        }

        vector<int> father(cnt);
        vector<double> weight(cnt, 1.0);

        for (int i = 0; i < cnt; ++i) {
            father[i] = i;
        }

        for (int i = 0; i < equations.size(); ++i) {
            int a = points[equations[i][0]];
            int b = points[equations[i][1]];
            double w = values[i];
            merge(a, b, w, father, weight);
        }

        vector<double> ret;
        for (auto item : queries) {
            if (points.find(item[0]) == points.end() or points.find(item[1]) == points.end()) {
                ret.push_back(-1.0);
                continue;
            }
            int a = points[item[0]];
            int b = points[item[1]];
            int fa = find(a, father, weight), fb = find(b, father, weight);
            if (fa != fb) {
                ret.push_back(-1.0);
            }
            else {
                ret.push_back(weight[a] / weight[b]);
            }
        }

        for (int i = 0; i < cnt; ++i) {
            printf("%d %lf \n", father[i], weight[i]);
        }

        return ret;
    }
};