Last active
January 26, 2022 19:11
-
-
Save mrchnk/aca471326e0de2b4e17ab3f226127b45 to your computer and use it in GitHub Desktop.
Segment Tree (update by index, query by segment)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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