3600. 升級後的最大生成樹穩定性

思路

1. 先把必須要被選的邊選進來，
   • 判斷有沒有形成環，一旦形成就不是最小生成樹了。
   • 同時將不一定需要選擇的邊記住。
   • 同時將權重列入參考答案，也就是跟res取最小值，
2. 從可選可不選的邊當中找權重最大的幾條邊加入到生成樹當中，這一步用 Kruskal。
3. 選完之後，把前 k 小的邊都強化，再跟res比較最小值。

提交成功在最後一分鐘，太刺激了。

程式碼

Kruskal

struct Edge {
    int u, v, w;
    Edge(int u, int v, int w) : u(u), v(v), w(w) {}
    bool operator<(const Edge& other) {
        return this->w < other.w;
    }
};
class Solution {
private:
    // 並查集
    static const int MX = 1e5+1;
    int parent[MX];
    int size[MX];
    int find(int x) {
        if(x != parent[x]) {
            parent[x] = find(parent[x]);
        }
        return parent[x];
    }
    bool isSameSet(int x, int y) {
        return find(x) == find(y);
    }
    void unite(int x, int y) {
        int rootX = find(x);
        int rootY = find(y);
        if(rootX != rootY) {
            if(size[rootX] < size[rootY]) {
                parent[rootX] = rootY;
                size[rootY] += size[rootX];
            }
            else {
                parent[rootY] = rootX;
                size[rootX] += size[rootY];
            }
        }
    }
public:
    int maxStability(int n, vector<vector<int>>& edges, int k) {
        iota(parent, parent + n, 0);
        fill(size, size + n, 1);

        int edgeCount = 0;
        int res = INT_MAX;
        vector<Edge> optEdges; // 紀錄可以選的邊
        optEdges.reserve(edges.size());
        for(int i = 0; i < edges.size(); i++) {
            auto& edge = edges[i];
            if(edge.back()) { // 如果是必須要被加入的邊，事先加入。
                if(isSameSet(edge[0], edge[1])) { // 已經在同一個集合當中，表示這是多餘的邊
                    return -1;
                }
                unite(edge[0], edge[1]);
                edgeCount++;
                res = min(res, edge[2]);
            }
            else { // 沒有必須要被加入的邊，紀錄到 pq 當中
                optEdges.push_back(Edge(edge[0], edge[1], edge[2]));
            }
        }
        if(edgeCount == n - 1) return res; // 只需要靠 must 的橋樑連接

        sort(optEdges.rbegin(), optEdges.rend()); // 大到小排序
        
        vector<int> chosenWeights; // 被選擇到生成樹當中，可以強化的邊權。
        for(int i = 0; i < optEdges.size(); i++) {
            Edge edge = optEdges[i];
            if(isSameSet(edge.u, edge.v)) continue;
            unite(edge.u, edge.v);
            chosenWeights.push_back(edge.w);
            if(++edgeCount == n - 1) break;
        }
        if(edgeCount != n - 1) return -1;
        int used = 0;
        for(int i = chosenWeights.size() - 1; i >= 0; i--) { // 從小到大遍歷
            if(used++ < k) {
                res = min(res, (chosenWeights[i] << 1));
            }
            else {
                res = min(res, chosenWeights[i]);
            }
        }
        return res;
    }
};

複雜度分析

• 時間複雜度：$O(m\log{m})$
• 空間複雜度：$O(n)$