meta hacker cup (very hard problem) - discovered linear solution!

d2_running_on_fumes.py

"""
Rusty Rob
https://x.com/robertkingNZ

For this hard problem I discovered that an innovative linear solution exists and thought it was worth sharing.

Problem statement
https://www.facebook.com/codingcompetitions/hacker-cup/2020/qualification-round/problems/D2

TLDR Problem statement:
Given a tree of gas stations, each station has a fill cost.
Whats the cheapest way from start node to end node, given the cars tank capacity.

Youtube video of solution:
https://youtu.be/I1YpqYFHKTM

TLDR solution:
Traverse tree from start to finish, collect gas stations in a monotonically increasing queue.
If we find a cheaper station, discard previous expensive stations.
the cost after the current station is equal to the cost of the first station in the queue plus the cost to fill at the current station.
if the first station in the queue is out of range (based on tank capacity), pop it from the queue
here's where it gets interesting:
if we reach a subtree, we can traverse down it and find the best cumulative fill cost for each level.
We can then "fold" this subtree back onto the original path. We can do this in a cost twice the height of the subtree.
to do this, we essentially merge together two monotonic queues in cost 2*H.
While it feels like its O(N^2), each node only has two associated operations, so the cost is O(N).
"""

from collections import deque
from heapq import merge
for case_num in range(1, int(input()) + 1):
    (N, M, A, B) = map(int, input().split())
    A, B = A - 1, B - 1
    g = [[] for _ in range(N)]
    costs = [0]*N
    def parse_input():
        for i in range(N):
            (p, c) = map(int, input().split())
            if i == A:
                c = 0
            costs[i] = c
            p = p - 1
            if p != -1:
                g[p].append(i)
                g[i].append(p)
    parse_input()
    def costs_below(node):
        fringe = [node]
        while True:
            new_fringe = []
            bc = 0
            for parent in fringe:
                for child in g[parent]:
                    if child not in visited:
                        visited.add(child)
                        new_fringe.append(child)
                        if bc == 0 or costs[child] < bc:
                            bc = costs[child]
            if not new_fringe:
                break
            yield bc
            fringe = new_fringe

    def get_path():
        stack = [A]
        visited = set(stack)
        m = {}
        while stack:
            parent = stack.pop()
            for child in g[parent]:
                if child not in visited:
                    m[child] = parent
                    if child == B:
                        path = [child]
                        while child != A:
                            child = m[child]
                            path.append(child)
                        return path[::-1]
                    visited.add(child)
                    stack.append(child)
    path = get_path()
    visited = set(path)
    q = deque([(0, 0)])
    ans = -1
    i = 0
    for i, node in enumerate(path):
        best_cost = q[0][1]
        if i == len(path) - 1:
            ans = best_cost
        if i - q[0][0] == M:
            q.popleft()
        c = costs[node]
        if c != 0:
            while q and q[-1][1] >= best_cost + c:
                q.pop()
            q.append((i, best_cost + c))
        elif not q:
            break
        k = 0
        better = deque()
        for j, cb in enumerate(costs_below(node), 1):
            best_cost = q[k][1]
            if i + j - q[k][0] == M:
                k += 1
            if cb != 0:
                new_cost = cb + best_cost
                if not better or new_cost < better[-1][1]:
                    better.appendleft((i - j, new_cost))
            if k == len(q) or (j + 1 == i - q[k][0]):
                break
        if better:
            q2 = deque()
            while q and q[-1][0] >= better[0][0]:
                q2.appendleft(q.pop())
            for vi, vc in merge(q2, better):
                while q and q[-1][1] >= vc:
                    q.pop()
                if not q or q[-1][0] != vi:
                    q.append((vi, vc))
    print('Case #{}: {}'.format(case_num, ans))