package main

import (
    "fmt"
    "sort"
)

func main() {
    graph := &Graph{}
    graph.nodes = make(map[int]*Node)
    graph.nodes[0] = &Node{}
    graph.nodes[1] = &Node{}
    graph.nodes[2] = &Node{}

    graph.edges = make(map[*Edge]struct{})
    graph.edges[&Edge{weight: 22, from: graph.nodes[0], to: graph.nodes[1]}] = struct{}{}
    graph.edges[&Edge{weight: 33, from: graph.nodes[1], to: graph.nodes[2]}] = struct{}{}
    graph.edges[&Edge{weight: 11, from: graph.nodes[2], to: graph.nodes[0]}] = struct{}{}
    ret := kruskalMST(graph)
    fmt.Println("结果：")
    for a, _ := range ret {
        fmt.Println(a.weight)
    }
}

type Edge struct {
    weight int
    from   *Node
    to     *Node
}

// 点结构的描述
type Node struct {
    value int
    in    int
    out   int
    nexts []*Node
    edges []*Edge
}
type Graph struct {
    nodes map[int]*Node
    edges map[*Edge]struct{}
}

func printPriorityQueue(priorityQueue []*Edge) {
    for _, edge := range priorityQueue {
        fmt.Println(edge.weight)
    }
}

func kruskalMST(graph *Graph) map[*Edge]struct{} {
    unionFind := &UnionFind{}
    unionFind.makeSets(graph.nodes)
    // 从小的边到大的边，依次弹出，小根堆！
    priorityQueue := make([]*Edge, 0)
    for edge, _ := range graph.edges {
        priorityQueue = append(priorityQueue, edge)
    }
    fmt.Println("排序前：")
    printPriorityQueue(priorityQueue)
    //排序
    sort.SliceStable(priorityQueue, func(i int, j int) bool {
        return priorityQueue[i].weight > priorityQueue[j].weight
    })
    fmt.Println("--------")
    fmt.Println("排序后：")
    printPriorityQueue(priorityQueue)
    fmt.Println("--------")
    result := make(map[*Edge]struct{})
    for len(priorityQueue) > 0 { // M 条边
        edge := priorityQueue[len(priorityQueue)-1]
        priorityQueue = priorityQueue[0 : len(priorityQueue)-1]
        if !unionFind.isSameSet(edge.from, edge.to) { // O(1)

            result[edge] = struct{}{}
            unionFind.union(edge.from, edge.to)
        }
    }
    return result
}

type UnionFind struct {
    // key 某一个节点， value key节点往上的节点
    fatherMap map[*Node]*Node
    // key 某一个集合的代表节点, value key所在集合的节点个数
    sizeMap map[*Node]int
}

func (this *UnionFind) makeSets(nodes map[int]*Node) {
    this.fatherMap = make(map[*Node]*Node)
    this.sizeMap = make(map[*Node]int)
    for _, node := range nodes {
        this.fatherMap[node] = node
        this.sizeMap[node] = 1
    }
}

func (this *UnionFind) findFather(n *Node) *Node {
    path := make([]*Node, 0)
    for n != this.fatherMap[n] {
        path = append(path, n)
        n = this.fatherMap[n]
    }
    for len(path) > 0 {

        this.fatherMap[path[len(path)-1]] = n
        path = path[0 : len(path)-1]
    }
    return n
}

func (this *UnionFind) isSameSet(a *Node, b *Node) bool {
    return this.findFather(a) == this.findFather(b)
}

func (this *UnionFind) union(a *Node, b *Node) {
    if a == nil || b == nil {
        return
    }
    aDai := this.findFather(a)
    bDai := this.findFather(b)
    if aDai != bDai {
        aSetSize := this.sizeMap[aDai]
        bSetSize := this.sizeMap[bDai]
        if aSetSize <= bSetSize {
            this.fatherMap[aDai] = bDai
            this.sizeMap[bDai] = aSetSize + bSetSize
            delete(this.sizeMap, aDai)
        } else {
            this.fatherMap[bDai] = aDai
            this.sizeMap[aDai] = aSetSize + bSetSize
            delete(this.sizeMap, bDai)
        }
    }
}