Skip to content

Instantly share code, notes, and snippets.

@appgurueu
Created April 30, 2023 12:26
Show Gist options
  • Save appgurueu/4d15978b283ce24147673cdc991ae3e4 to your computer and use it in GitHub Desktop.
Save appgurueu/4d15978b283ce24147673cdc991ae3e4 to your computer and use it in GitHub Desktop.
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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment