动态规划--图的最短路径问题(go实现)

超过 100%人提交,性能不错

package leetcode

import (
	"container/heap"
	"errors"
)

type graph struct {
	vertex []*listNode
	count  int
}

type listNode struct {
	to       int
	priority int
	next     *listNode
}

func newGraph(times [][]int, n int) (*graph, error) {
	if n == 0 {
		return nil, errors.New("n is require not zore")
	}
	g := &graph{
		vertex: make([]*listNode, n, n),
		count:  n,
	}

	for _, time := range times {
		if len(time) < 3 {
			return nil, errors.New("len(time) must be more than 3")
		}

		node := &listNode{
			to:       time[1] - 1,
			priority: time[2],
		}
		if g.vertex[time[0]-1] == nil {
			g.vertex[time[0]-1] = node
		} else {
			node.next = g.vertex[time[0]-1]
			g.vertex[time[0]-1] = node
		}
	}
	return g, nil
}

// [[2,1,1],[2,3,1],[3,4,1]], n = 4, k = 2
func networkDelayTime(times [][]int, n int, k int) int {
	k--
	graph, _ := newGraph(times, n)
	inqueue := make([]bool, n, n)
	preVertex := make([]int, n, n)
	entryList := make([]*entry, n, n)
	for i := 0; i < n; i++ {
		entryList[i] = &entry{
			vertex:   i,
			priority: math.MaxInt,
		}
	}

	pq := &PQ{}
	var vertex *listNode
	inqueue[k] = true
	entryList[k].priority = 0
	preEntry := &entry{
		vertex:   k,
		priority: 0,
	}
	heap.Push(pq, preEntry)

	for pq.Len() != 0 {
		e := heap.Pop(pq)
		preEntry = e.(*entry)
		for vertex = graph.vertex[preEntry.vertex]; vertex != nil; vertex = vertex.next {
			if inqueue[vertex.to] && preEntry.priority+vertex.priority < entryList[vertex.to].priority {
				entryList[vertex.to].priority = preEntry.priority + vertex.priority
				preVertex[vertex.to] = preEntry.vertex
				pq.update(entryList[vertex.to])
			} else if !inqueue[vertex.to] {
				e1 := &entry{
					vertex:   vertex.to,
					priority: preEntry.priority + vertex.priority,
				}

				inqueue[vertex.to] = true
				entryList[vertex.to] = e1
				preVertex[vertex.to] = preEntry.vertex
				heap.Push(pq, e1)
			}
		}
	}

	for _, e := range inqueue {
		if !e {
			return -1
		}
	}

	return entryList[preEntry.vertex].priority
}

// entry 是 priorityQueue 中的元素
type entry struct {
	vertex   int
	priority int
	// index 是 entry 在 heap 中的索引号
	// entry 加入 Priority Queue 后， Priority 会变化时，很有用
	// 如果 entry.priority 一直不变的话，可以删除 index
	index int
}

// PQ implements heap.Interface and holds entries.
type PQ []*entry

func (pq PQ) Len() int { return len(pq) }

func (pq PQ) Less(i, j int) bool {
	return pq[i].priority < pq[j].priority
}

func (pq PQ) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

// Push 往 pq 中放 entry
func (pq *PQ) Push(x interface{}) {
	temp := x.(*entry)
	temp.index = len(*pq)
	*pq = append(*pq, temp)
}

// Pop 从 pq 中取出最优先的 entry
func (pq *PQ) Pop() interface{} {
	temp := (*pq)[len(*pq)-1]
	temp.index = -1 // for safety
	*pq = (*pq)[0 : len(*pq)-1]
	return temp
}

// update modifies the priority and value of an entry in the queue.
func (pq *PQ) update(entry *entry) {
	heap.Fix(pq, entry.index)
}