https://leetcode.com/problems/checking-existence-of-edge-length-limited-paths/description/

pufs.go

import "sort"

// This is not particularly efficient.
// A more efficient implementation might either be B-Tree like,
// might eliminate the index, might distinguish parents and leaves,
// and might strip parents of their values;
// to alleviate the mem management overhead, I'd wish
// for an arena allocator, but this isn't Zig

type PSlice[T any] struct {
	i int
	v T
	left, right *PSlice[T]
}

func NewPSlice[T any](from, to int, init func(i int) T) *PSlice[T] {
	if to <= from {
		return nil
	}
	mid := (from + to) >> 1
	return &PSlice[T]{
		mid,
		init(mid),
		NewPSlice[T](from, mid, init), // mid is exclusive
		NewPSlice[T](mid + 1, to, init),
	}
}

func (self *PSlice[T]) set(i int, v T) *PSlice[T] {
    if self == nil {
    	panic("out of bounds PSlice access")
   	}
   	cp := *self
   	if i < self.i {
   		cp.left = self.left.set(i, v)
   	} else if i > self.i {
   		cp.right = self.right.set(i, v)
   	} else {
   		cp.v = v
   	}
	return &cp
}

func (self *PSlice[T]) get(i int) T {
	if self == nil {
		panic("out of bounds PSlice access")
	}
	if i == self.i {
		return self.v
	}
	if i < self.i {
		return self.left.get(i)
	}
	return self.right.get(i)
}

type PSetNode struct {
	v int // HACK unused if not a leaf
	left, right *PSetNode
}

type PSet struct {
	count int
	root *PSetNode
}

// Persistent Union Find
type SetIdx = int
type PUF struct {
	setidx *PSlice[SetIdx]
	sets *PSlice[PSet]
}

func NewPUF(size int) *PUF {
	setidx := NewPSlice[SetIdx](0, size, func(i int) SetIdx {
		return i
	})
	sets := NewPSlice[PSet](0, size, func(i int) PSet {
		return PSet{1, &PSetNode{i, nil, nil}}
	})
	return &PUF{setidx, sets}
}

func (self *PUF) SameSet(i, j int) bool {
	return self.setidx.get(i) == self.setidx.get(j)
}

func (self *PUF) Union(i, j int) *PUF {
	setidx := self.setidx
	isi, jsi := setidx.get(i), setidx.get(j)
	if isi == jsi {
		return self
	}
	sets := self.sets
	is, js := sets.get(isi), sets.get(jsi)
	if is.count < js.count {
		// enforce is.count >= js.count
		isi, jsi = jsi, isi
		is, js = js, is 
	}
	us := PSet{is.count + js.count, &PSetNode{-1, is.root, js.root}}
	sets = sets.set(isi, us)
	var iter func(node *PSetNode)
	iter = func(node *PSetNode) {
		if node == nil {
			return
		}
		iter(node.left)
		iter(node.right)
		if node.left == nil && node.right == nil {
			// TODO I don't really like this since it is n log n - the logs are stacking!
			setidx = setidx.set(node.v, isi)
		}
	}
	iter(js.root)
	// technically we could remove jsi from self.sets
	// but that would complicate matters and we don't really care
	return &PUF{setidx, sets}
}

func distanceLimitedPathsExist(n int, edgeList [][]int, queries [][]int) []bool {
	pufs := make([]*PUF, len(edgeList) + 1)
	sort.Slice(edgeList, func(i, j int) bool {
		return edgeList[i][2] < edgeList[j][2]
	})
	pufs[0] = NewPUF(n) // index shifted by one for a reason
	for i, edge := range edgeList {
		from, to := edge[0], edge[1]
		pufs[i + 1] = pufs[i].Union(from, to)
	}
	res := make([]bool, len(queries))
	for i, query := range queries {
		from, to, limit := query[0], query[1], query[3]
		// Find index of first edge above limit
		j := sort.Search(len(pufs), func(k int) bool {
			return edgeList[k][3] >= limit
		})
		// Answer same set query in corresponding PUF *before* limit
		res[i] = pufs[j].SameSet(from, to)
	}
	return res
}