import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
* Prim算法
*
* 1).输入:一个加权连通图,其中顶点集合为V,边集合为E;
* 2).初始化:Vnew = {x},其中x为集合V中的任一节点(起始点),Enew = {},为空;
* 3).重复下列操作,直到Vnew = V:
* a.在集合E中选取权值最小的边<u, v>,其中u为集合Vnew中的元素,而v不在Vnew集合当中,并且v∈V(如果存在有多条满足前述条件即具有相同权值的边,则可任意选取其中之一);
* b.将v加入集合Vnew中,将<u, v>边加入集合Enew中;
* 4).输出:使用集合Vnew和Enew来描述所得到的最小生成树。
*/
public class Prim {

public static Set<Edge> primMST(Graph graph) {
// 解锁的边进入小根堆
PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());

// 哪些点被解锁出来了
HashSet<Node> nodeSet = new HashSet<>();
Set<Edge> result = new HashSet<>(); // 依次挑选的的边在result里
for (Node node : graph.nodes.values()) { // 随便挑了一个点
// node 是开始点
if (!nodeSet.contains(node)) {
nodeSet.add(node);
for (Edge edge : node.edges) { // 由一个点,解锁所有相连的边
priorityQueue.add(edge);
}
while (!priorityQueue.isEmpty()) {
Edge edge = priorityQueue.poll(); // 弹出解锁的边中,最小的边
Node toNode = edge.to; // 可能的一个新的点
if (!nodeSet.contains(toNode)) { // 不含有的时候,就是新的点
nodeSet.add(toNode);
result.add(edge);
for (Edge nextEdge : toNode.edges) {
priorityQueue.add(nextEdge);
}
}
}
}
// break;
}
return result;
}

public static class EdgeComparator implements Comparator<Edge> {

@Override
public int compare(Edge o1, Edge o2) {
return o1.weight - o2.weight;
}

}

// 请保证graph是连通图
// graph[i][j]表示点i到点j的距离,如果是系统最大值代表无路
// 返回值是最小连通图的路径之和
public static int prim(int[][] graph) {
int size = graph.length;
int[] distances = new int[size];
boolean[] visit = new boolean[size];
visit[0] = true;
for (int i = 0; i < size; i++) {
distances[i] = graph[0][i];
}
int sum = 0;
for (int i = 1; i < size; i++) {
int minPath = Integer.MAX_VALUE;
int minIndex = -1;
for (int j = 0; j < size; j++) {
if (!visit[j] && distances[j] < minPath) {
minPath = distances[j];
minIndex = j;
}
}
if (minIndex == -1) {
return sum;
}
visit[minIndex] = true;
sum += minPath;
for (int j = 0; j < size; j++) {
if (!visit[j] && distances[j] > graph[minIndex][j]) {
distances[j] = graph[minIndex][j];
}
}
}
return sum;
}

class Graph {

public HashMap<Integer, Node> nodes;

public HashSet<Edge> edges;

public Graph() {
nodes = new HashMap<>();
edges = new HashSet<>();
}

}

class Node {

public int value;

public int in;

public int out;

public ArrayList<Node> nexts;

public ArrayList<Edge> edges;

public Node(int value) {
this.value = value;
nexts = new ArrayList<>();
edges = new ArrayList<>();
}

}

class Edge {

// 权重
public int weight;

public Node from;

public Node to;

public Edge(int weight, Node from, Node to) {
this.weight = weight;
this.from = from;
this.to = to;
}

}
}

/* 如有意见或建议,欢迎评论区留言;如发现代码有误,欢迎批评指正 */