# Priority queue, whose elements are ints, and whose priorities can be
# adjusted without taking them out of the queue. Modeled after
# Sedgewick & Wayne's heap class, specifically IndexMinPQ.java,
# available at https://algs4.cs.princeton.edu/
#
# Jesper Larsson, Malmö University 2018–2020

from clo_decorators import clo_method

class VertexPQ:
    def __init__(self, max_n):
        n = 0
        pq = [-1] * (max_n+1)             # vertices in heap order
        qp = [-1] * max_n                 # positions of vertices in pq
        dist = [float('inf')] * max_n     # keys of vertices, lower dist means higher prio

        def exch(i, j):
            v = pq[i]
            w = pq[j]
            pq[i] = w
            pq[j] = v
            qp[w] = i
            qp[v] = j

        def swim(k):
            while k > 1 and dist[pq[k//2]] > dist[pq[k]]:
                exch(k, k//2)
                k //= 2

        def sink(k):
            nonlocal n
            while 2*k <= n:
                j = 2*k
                if j < n and dist[pq[j]] > dist[pq[j+1]]:
                    j += 1
                if dist[pq[k]] <= dist[pq[j]]:
                    break
                exch(k, j)
                k = j

        @clo_method(self)
        def set_dist(v, d):
            if qp[v] < 0:                 # vertex not in queue, add it
                nonlocal n
                n += 1
                pq[n] = v
                qp[v] = n
                dist[v] = d
                swim(n)
            else:                         # already in queue, modify
                grows = d > dist[v]
                dist[v] = d
                if grows:
                    sink(qp[v])
                else:
                    swim(qp[v])

        @clo_method(self)
        def get_dist(v):
            return dist[v]

        @clo_method(self)
        def del_min():
            nonlocal n
            min = pq[1]
            exch(1, n)
            n -= 1
            sink(1)
            qp[min] = -1
            return min

        @clo_method(self)
        def is_empty():
            return n == 0
        
    def __bool__(self):
        return not self.is_empty()