golang-graph

graph（比较乱只做参考）

// graph
package graph

import (
	"fmt"
	"sort"
	"sync"
)

// Adjacency matrix
//type adjGraph struct {
//	node []int
//	edges [][]int
//}

// Adjacency list
type Node struct {
	value int
}

type Edge struct {
	parent *Node
	child  *Node
	weight int
}

type UnionSet struct {
	parent map[Node]Node
	rank   map[Node]int
}

type edgeQueue struct {
	edges []Edge
}

func (n *Node) String() string {
	return fmt.Sprintf("%v", n.value)
}

type ListGraph struct {
	nodes []*Node
	edges map[Node][]Edge
	weights []Edge
	lock sync.RWMutex
}

const INT_MAX = int(^uint(0) >> 1)

func (g *ListGraph) AddNode(n *Node) {
	g.lock.Lock()
	defer g.lock.Unlock()
	g.nodes = append(g.nodes, n)
}

func (g *ListGraph) AddEdge(u, v *Node, w int) {
	g.lock.Lock()
	defer g.lock.Unlock()
	if g.edges == nil {
		g.edges = make(map[Node][]Edge)
	}
	g.edges[*u] = append(g.edges[*u], Edge{u, v, w})
	g.edges[*v] = append(g.edges[*v], Edge{v, u, w})
}

func (g *ListGraph) AddWeight(u, v *Node, w int) {
	g.lock.Lock()
	defer g.lock.Unlock()
	g.weights = append(g.weights, Edge{u, v, w})
}

func (g *ListGraph) String() {
	g.lock.Lock()
	defer g.lock.Unlock()
	str := ""
	for _, iNode := range g.nodes {
		str += iNode.String() + " -> "
		nexts := g.edges[*iNode]
		for _, next := range nexts {
			str += next.child.String() + " "
		}
		str += "\n"
	}
	fmt.Println(str)
}

// 实现heap接口
func (q edgeQueue) Len() int {
	return len(q.edges)
}

func (q edgeQueue) Swap(i, j int) {
	q.edges[i], q.edges[j] = q.edges[j], q.edges[i]
}

func (q edgeQueue) Less(i, j int) bool {
	return q.edges[i].weight < q.edges[j].weight
}

func (q *edgeQueue) Push(x interface{}) {
	q.edges = append(q.edges, x.(Edge))
}

func (q *edgeQueue) Pop() interface{} {
	x := q.edges[q.Len()-1]
	q.edges = q.edges[:q.Len()-1]
	return x
}

// Prim MST
func (g *ListGraph)Prim() []Edge {
	var edges []Edge
	var res []Edge
	visited := make(map[*Node]bool)
	n := len(g.nodes)
	if n == 0 {
		return nil
	}

	init := g.nodes[0]
	visited[init] = true
	edges = append(edges, g.edges[*init]...)

	for len(res) < n-1 {
		sort.Slice(edges, func(i, j int) bool {
			return edges[i].weight < edges[j].weight
		})

		for i := 0; i < len(edges); i++ {
			if !visited[edges[i].child] {
				//fmt.Printf("%dth: %v\n", i, edges)
				res = append(res, edges[i])
				visited[edges[i].child] = true
				edges = append(edges, g.edges[*(edges[i].child)]...)
				sort.Slice(edges, func(i, j int) bool {
					return edges[i].weight < edges[j].weight
				})
			}
		}
	}
	return res
}

//并查集
func NewUnionSet(nodes []*Node) UnionSet {
	set, rank := make(map[Node]Node, len(nodes)), make(map[Node]int, len(nodes))
	for _, node := range nodes {
		set[*node] = *node
		rank[*node] = 1
	}
	return UnionSet{set, rank}
}

func (u UnionSet) Find(node Node) Node {
	for ; u.parent[node] != node; node = u.parent[node] {}
	return node
}

func (u UnionSet) Union(x, y Node) {
	fx, fy := u.Find(x), u.Find(y)
	if fx == fy {
		return
	}

	if u.rank[fx] < u.rank[fy] {
		u.parent[fx] = fy
		return
	}
	u.parent[fy] = fx
	if u.rank[fx] == u.rank[fy] {
		u.rank[fx]++
	}
}

//Kruskal MST
func (g *ListGraph) Kruskal() []Edge {
	var edges, res []Edge
	for _, weight := range g.weights {
		edges = append(edges, weight)
	}

	sort.Slice(edges, func(i, j int) bool {
		return edges[i].weight < edges[j].weight
	})

	unionSet := NewUnionSet(g.nodes)
	for _, e := range edges {
		if n1, n2 := unionSet.Find(*(e.parent)), unionSet.Find(*(e.child)); n1 != n2 {
			fmt.Println("unionSet:", unionSet)
			unionSet.Union(*(e.parent), *(e.child))
			res = append(res, e)
		}
	}
	fmt.Println("unionSet:", unionSet)
	return res
}

// BFS
func (g *ListGraph) BFS(f func(*Node)) {
	g.lock.RLock()
	defer g.lock.RUnlock()
	visited := make(map[*Node]bool)
	fmt.Println(visited)
	q := []*Node{}

	n := g.nodes[0]
	visited[n] = true
	q = append(q, n)

	for len(q) > 0 {
		var current *Node
		current, q = q[0], q[1:]
		for _, edge := range g.edges[*current] {
			if !visited[edge.child] {
				q = append(q, edge.child)
				visited[edge.child] = true
			}
		}
		if f != nil {
			f(current)
		}
	}
}



// DFS
func (g *ListGraph) DFS(f func(*Node)) {
	visited := make(map[*Node]bool)
	node := g.nodes[0]
	g.dfsRecursive(node, visited, f)
}

func (g *ListGraph) dfsRecursive(startingNode *Node, visited map[*Node]bool, f func(*Node)) {
	if visited[startingNode] == true {
		return
	}
	visited[startingNode] = true
	f(startingNode)
	for _, edge := range g.edges[*startingNode] {
		g.dfsRecursive(edge.child, visited, f)
	}

}

//Dijkstra 单源最短路径
func (g *ListGraph) Dijkstra(Id int) ([]Node, int) {
	s, n := g.nodes[Id], len(g.nodes)
	adj, visited, pre := make(map[*Node]int), make(map[*Node]bool), make(map[*Node]Node)

	for i := 0; i < n; i++ {
		if s == g.nodes[i] {
			adj[s] = 0
			visited[s] = true
		}else {
			adj[g.nodes[i]] = INT_MAX/2
			visited[g.nodes[i]] = false
		}
		pre[g.nodes[i]] = *(g.nodes[i])
	}
	for _, v := range g.edges[*s] {
		adj[v.child] = v.weight
	}
	fmt.Println("adj:", adj)
	for i := 0; i < n; i++ {
		tmpWeight := INT_MAX
		for node, w := range adj {
			if !visited[node] && w < tmpWeight {
				s = node
				tmpWeight = w
			}
		}
		for _, v := range g.edges[*s] {
			if adj[v.child] > adj[v.parent] + adj[v.child] {
				adj[v.child] = adj[v.parent] + adj[v.child]
				pre[v.child] = *(v.parent)
			}
		}
	}

	var res []Node
	for i := 0; i < n; i++ {
		res = append(res, pre[g.nodes[i]])
	}
	fmt.Println("adj:", adj)
	return res, adj[g.nodes[n-1]]
}

//func maxInt(a, b int) int {
//	if a > b {
//		return b
//	} else {
//		return a
//	}
//}

// graph_test
package graph

import (
	"fmt"
	"testing"
)

//func TestListGraph_AddNode(t *testing.T) {
//	g := listGraph{}
//	nodes := []*Node{
//		{1},
//		{2},
//		{3},
//		{4},
//		{5},
//	}
//	for _, v := range nodes {
//		g.AddNode(v)
//	}
//
//	g.AddEdge(nodes[0], nodes[1])
//	g.AddEdge(nodes[0], nodes[4])
//	g.AddEdge(nodes[1], nodes[2])
//	g.AddEdge(nodes[1], nodes[3])
//	g.AddEdge(nodes[1], nodes[4])
//	g.AddEdge(nodes[2], nodes[3])
//	g.AddEdge(nodes[3], nodes[4])
//	g.String()
//}

func TestListGraph_BFS(t *testing.T) {
	g := ListGraph{}
	nodes := []*Node{
		{1},
		{2},
		{3},
		{4},
		{5},
		{6},
	}
	for _, v := range nodes {
		g.AddNode(v)
	}

	g.AddEdge(nodes[0], nodes[1], 1)
	g.AddEdge(nodes[0], nodes[2], 12)
	g.AddEdge(nodes[1], nodes[2], 9)
	g.AddEdge(nodes[1], nodes[3], 3)
	g.AddEdge(nodes[2], nodes[4], 5)
	g.AddEdge(nodes[3], nodes[2], 4)
	g.AddEdge(nodes[3], nodes[4], 13)
	g.AddEdge(nodes[3], nodes[5], 15)
	g.AddEdge(nodes[4], nodes[5], 4)

	g.AddWeight(nodes[0], nodes[1], 1)
	g.AddWeight(nodes[0], nodes[2], 12)
	g.AddWeight(nodes[1], nodes[2], 9)
	g.AddWeight(nodes[1], nodes[3], 3)
	g.AddWeight(nodes[2], nodes[4], 5)
	g.AddWeight(nodes[3], nodes[2], 4)
	g.AddWeight(nodes[3], nodes[4], 13)
	g.AddWeight(nodes[3], nodes[5], 15)
	g.AddWeight(nodes[4], nodes[5], 4)

	g.String()

	f := func(n *Node){
		fmt.Printf("%v ", n.String())
	}

	g.BFS(f)
	fmt.Println("")
	g.DFS(f)
	fmt.Println("")

	MSTK := g.Kruskal()
	for _, e := range MSTK {
		fmt.Println(*(e.parent), *(e.child), e.weight)
	}
	fmt.Println("================")
	//MSTP := g.Prim()
	//for _, e := range MSTP {
	//	fmt.Println(*(e.parent), *(e.child), e.weight)
	//}

	minPath, w := g.Dijkstra(0)
	fmt.Println("path:", minPath)
	fmt.Println(w)
}