Weighted Union-Find Potential Query

Problem

Maintain a weighted union-find for difference constraints. operations contains ["union", a, b, w] meaning val[b] - val[a] = w, or ["query", a, b]. A union records True when a and b were in different components (and adds the relation), False when they were already connected; a query records val[b] - val[a] when connected, otherwise None. Use path compression that accumulates potentials.

Hint

The potential pot[x] means val[x] - val[parent[x]]; when attaching rb under ra, set pot[rb] = w + pot[a] - pot[b].

Test cases

case 1

in: [[["union",0,1,5],["query",0,1],["union",1,2,3],["query",0,2],["query",2,0],["union",0,2,8],["query",1,2],["query",0,3]]]

expected: [true,5,true,8,-8,false,3,null]

case 2

in: [[["query",5,5],["union",5,6,10],["query",6,5],["union",6,7,4],["query",5,7]]]

expected: [0,true,-10,true,14]

case 3

in: [[["union",0,1,1],["union",1,2,2],["query",0,2],["union",0,2,99],["query",1,1],["query",2,1]]]

expected: [true,true,3,false,0,-2]

Starter codePython

def weighted_dsu(operations):
    # Your code here
    pass

3 test cases · run your code against them in the editor