Skip to content

Instantly share code, notes, and snippets.

@mrchnk
Last active January 26, 2022 19:11
Show Gist options
  • Save mrchnk/aca471326e0de2b4e17ab3f226127b45 to your computer and use it in GitHub Desktop.
Save mrchnk/aca471326e0de2b4e17ab3f226127b45 to your computer and use it in GitHub Desktop.
Segment Tree (update by index, query by segment)
function SegmentTree(size, merge, zero)
local tree = {}
local function query(v, tl, tr, l, r)
if l == tl and r == tr then
return tree[v] or zero
end
local tm = math.floor((tl + tr) / 2)
if (r <= tm) then
return query(v * 2, tl, tm, l, r)
elseif (l >= tm + 1) then
return query(v * 2 + 1, tm + 1, tr, l, r)
else
local lq = query(v * 2, tl, tm, l, tm)
local rq = query(v * 2 + 1, tm + 1, tr, tm + 1, r)
return merge(lq, rq)
end
end
local function update(v, tl, tr, index, value)
if tl == tr then
tree[v] = value
return value
else
local tm = math.floor((tl + tr) / 2)
local vl, vr
if index <= tm then
vl = update(v * 2, tl, tm, index, value)
vr = tree[v * 2 + 1] or zero
else
vl = tree[v * 2] or zero
vr = update(v * 2 + 1, tm + 1, tr, index, value)
end
value = merge(vl, vr)
tree[v] = value
return value
end
end
return {
tree = tree,
query = function (l, r)
return query(1, 1, size, l, r)
end,
update = function (index, value)
return update(1, 1, size, index, value)
end
}
end
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment