Leetcode Problem 787. "Cheapest Flights Within K Stops" Solution

leetcode_787.go

// Solution to https://leetcode.com/problems/cheapest-flights-within-k-stops/description/
// probably not the intended solution; uses an abridged Dijkstra to solve the much more general problem
// of finding all shortest paths to a given node using exactly k nodes

package main

import (
  "math"
  "container/heap"
)

type OutgoingFlight struct {
    to, distance int
}

type KNode struct {
    node, k, distance, idx int
}

// A PriorityQueue implements heap.Interface and holds KNodes.
type PriorityQueue []*KNode

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

func (pq PriorityQueue) Less(i, j int) bool {
    if pq[i].k == pq[j].k {
        return pq[i].distance < pq[j].distance
    }
	return pq[i].k < pq[j].k
}

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].idx = i
	pq[j].idx = j
}

func (pq *PriorityQueue) Push(x interface{}) {
	n := len(*pq)
	knode := x.(*KNode)
	knode.idx = n
	*pq = append(*pq, knode)
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	knode := old[n-1]
	old[n-1] = nil  // avoid memory leak
	knode.idx = -1 // for safety
	*pq = old[0 : n-1]
	return knode
}

func (pq *PriorityQueue) update(knode *KNode, distance int) {
	knode.distance = distance
	heap.Fix(pq, knode.idx)
}


func findCheapestPrice(n int, flights [][]int, src int, dst int, k int) int {
    outgoingFlights := make([][]OutgoingFlight, n)
    for _, flight := range flights {
        from, to, price := flight[0], flight[1], flight[2]
        outgoingFlights[from] = append(outgoingFlights[from], OutgoingFlight{to, price})
    }
    
    k++ // length of path = #stops + 1
    knodes := make([][]*KNode, n)
    for i := range knodes {
        knodes[i] = make([]*KNode, k+1)
        for j := range knodes[i] {
            knodes[i][j] = &KNode{
                node: i,
                k: j,
                distance: math.MaxInt,
                idx: -1,
            }
        }
    }
    knodes[src][0].distance = 0

    queue := make(PriorityQueue, 1)
    queue[0] = knodes[src][0]
    heap.Init(&queue)
    for len(queue) > 0 {
        knode := heap.Pop(&queue).(*KNode)
        fmt.Println(knode)
        if knode.k >= k {
            continue
        }
        for _, outFlight := range outgoingFlights[knode.node] {
            tryDist := knode.distance + outFlight.distance
            dstKnode := knodes[outFlight.to][knode.k + 1]
            if dstKnode.idx == -1 {
                dstKnode.distance = tryDist
                heap.Push(&queue, dstKnode)
            } else if tryDist < dstKnode.distance {
                queue.update(dstKnode, tryDist)
            }
        }
    }

    minDist := math.MaxInt
    for _, knode := range knodes[dst] {
        if knode.distance < minDist {
            minDist = knode.distance
        }
    }
    if minDist == math.MaxInt {
        return -1
    }
    return minDist
}