请给出一段k-MST代码

以下是一个使用Kruskal算法求解最小生成树的代码示例：

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
        
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def union(self, x, y):
        px, py = self.find(x), self.find(y)
        if px == py:
            return False
        if self.rank[px] > self.rank[py]:
            self.parent[py] = px
        elif self.rank[px] < self.rank[py]:
            self.parent[px] = py
        else:
            self.parent[py] = px
            self.rank[px] += 1
        return True

def kruskal(n, edges, k):
    uf = UnionFind(n)
    edges.sort(key=lambda x: x[2])
    mst = []
    for e in edges:
        if uf.union(e[0], e[1]):
            mst.append(e)
            if len(mst) == k:
                break
    return mst

这里 `n` 表示节点数量， `edges` 是一个由三元组 `(u, v, w)` 表示的边列表，其中 `(u, v)` 是边的两个端点， `w` 是边的权重。 `k` 表示需要找到的最小生成树的边数。函数返回一个包含 `k` 条边的最小生成树。