package main

import (
	"fmt"
	"slices"
)

// 完整的並查集模板，見我的數據結構題單
type unionFind struct {
	fa []int // 代表元
	cc int   // 連通塊個數
}

func newUnionFind(n int) unionFind {
	fa := make([]int, n)
	// 一開始有 n 個集合 {0}, {1}, ..., {n-1}
	// 集合 i 的代表元是自己
	for i := range fa {
		fa[i] = i
	}
	return unionFind{fa, n}
}

// 返回 x 所在集合的代表元
// 同時做路徑壓縮，也就是把 x 所在集合中的所有元素的 fa 都改成代表元
func (u unionFind) find(x int) int {
	// 如果 fa[x] == x，則表示 x 是代表元
	if u.fa[x] != x {
		u.fa[x] = u.find(u.fa[x]) // fa 改成代表元
	}
	return u.fa[x]
}

// 把 from 所在集合合併到 to 所在集合中
func (u *unionFind) merge(from, to int) {
	x, y := u.find(from), u.find(to)
	if x == y { // from 和 to 在同一個集合，不做合併
		return
	}
	u.fa[x] = y // 合併集合。修改後就可以認為 from 和 to 在同一個集合了
	u.cc--      // 成功合併，連通塊個數減一
}

func minCost(n int, edges [][]int, k int) int {
	if k == n {
		return 0
	}

	slices.SortFunc(edges, func(a, b []int) int { return a[2] - b[2] })
	u := newUnionFind(n)
	for _, e := range edges {
		u.merge(e[0], e[1])
		if u.cc == k {
			return e[2]
		}
	}
	panic("impossible")
}

func main() {
	n := 5
	edges := [][]int{{0, 1, 4}, {1, 2, 3}, {1, 3, 2}, {3, 4, 6}}
	k := 2
	result := minCost(n, edges, k)
	fmt.Println(result)
}

# -*-coding:utf-8-*-

from typing import List

class UnionFind:
    def __init__(self, n: int):
        self.fa = list(range(n))  # 代表元
        self.cc = n               # 連通塊個數
    
    def find(self, x: int) -> int:
        """返回 x 所在集合的代表元，同時做路徑壓縮"""
        if self.fa[x] != x:
            self.fa[x] = self.find(self.fa[x])
        return self.fa[x]
    
    def merge(self, from_: int, to: int) -> None:
        """把 from_ 所在集合合併到 to 所在集合中"""
        x, y = self.find(from_), self.find(to)
        if x == y:  # from_ 和 to 在同一個集合，不做合併
            return
        self.fa[x] = y  # 合併集合
        self.cc -= 1    # 成功合併，連通塊個數減一

def minCost(n: int, edges: List[List[int]], k: int) -> int:
    if k == n:
        return 0
    
    # 按邊權從小到大排序
    edges.sort(key=lambda e: e[2])
    uf = UnionFind(n)
    
    for e in edges:
        uf.merge(e[0], e[1])
        if uf.cc == k:
            return e[2]
    
    raise ValueError("impossible")

if __name__ == "__main__":
    n = 5
    edges = [[0, 1, 4], [1, 2, 3], [1, 3, 2], [3, 4, 6]]
    k = 2
    result = minCost(n, edges, k)
    print(result)

#include <iostream>
#include <vector>
#include <algorithm>
#include <stdexcept>

using namespace std;

class UnionFind {
private:
    vector<int> fa; // 代表元
    int cc;         // 連通塊個數

public:
    // 構造函數
    UnionFind(int n) : fa(n), cc(n) {
        // 一開始有 n 個集合 {0}, {1}, ..., {n-1}
        // 集合 i 的代表元是自己
        for (int i = 0; i < n; i++) {
            fa[i] = i;
        }
    }

    // 返回 x 所在集合的代表元，同時做路徑壓縮
    int find(int x) {
        if (fa[x] != x) {
            fa[x] = find(fa[x]); // fa 改成代表元
        }
        return fa[x];
    }

    // 把 from 所在集合合併到 to 所在集合中
    void merge(int from, int to) {
        int x = find(from);
        int y = find(to);
        if (x == y) { // from 和 to 在同一個集合，不做合併
            return;
        }
        fa[x] = y; // 合併集合
        cc--;      // 成功合併，連通塊個數減一
    }

    // 獲取當前連通塊個數
    int getCC() const {
        return cc;
    }
};

int minCost(int n, vector<vector<int>>& edges, int k) {
    if (k == n) {
        return 0;
    }

    // 按邊權從小到大排序
    sort(edges.begin(), edges.end(),
         [](const vector<int>& a, const vector<int>& b) {
             return a[2] < b[2];
         });

    UnionFind uf(n);
    for (const auto& e : edges) {
        uf.merge(e[0], e[1]);
        if (uf.getCC() == k) {
            return e[2];
        }
    }

    throw runtime_error("impossible");
}

int main() {
    int n = 5;
    vector<vector<int>> edges = {{0, 1, 4}, {1, 2, 3}, {1, 3, 2}, {3, 4, 6}};
    int k = 2;

    int result = minCost(n, edges, k);
    cout << result << endl;

    return 0;
}