[
segment tree
]
CodeForces Segment Trees, Part 1, 2.4
Problem statement
https://codeforces.com/edu/course/2/lesson/4/2/practice/contest/273278/problem/D
Solution
This segment tree will support two operations:
- Set element
i
equal tov
. - First element in tree starting from index
l
, which is>= x
.
To solve this, we will keep tree with maximums, then on each step we decide to goto the left or to the right.
Complexity
Standart compexities for segment tree.
Code
class SegmentTree:
def __init__(self, n, arr):
self.size = 1
while self.size < n:
self.size *= 2
self.T = [0] * (2 * self.size - 1)
self.arr = arr
def _build(self, x, lx, rx):
if rx - lx == 1:
if lx < len(self.arr):
self.T[x] = self.arr[lx]
else:
mx = (lx + rx) // 2
self._build(2 * x + 1, lx, mx)
self._build(2 * x + 2, mx, rx)
self.T[x] = max(self.T[2 * x + 1], self.T[2 * x + 2])
def build(self):
self._build(0, 0, self.size)
def _set(self, i, v, x, lx, rx):
if rx - lx == 1:
self.T[x] = v
return
mx = (lx + rx) // 2
if i < mx:
self._set(i, v, 2 * x + 1, lx, mx)
else:
self._set(i, v, 2 * x + 2, mx, rx)
self.T[x] = max(self.T[2 * x + 1], self.T[2 * x + 2])
def set(self, i, v):
self._set(i, v, 0, 0, self.size)
def first_above_(self, v, l, x, lx, rx):
if self.T[x] < v or rx <= l: # here we add one more bad case when we return -1
return -1
if rx - lx == 1:
return lx
mx = (lx + rx) // 2
res = self.first_above_(v, l, 2 * x + 1, lx, mx)
if res == -1:
res = self.first_above_(v, l, 2 * x + 2, mx, rx)
return res
def first_above(self, v, l):
return self.first_above_(v, l, 0, 0, self.size)
import io, os, sys
input = io.BytesIO(os.read(0,os.fstat(0).st_size)).readline
if __name__ == '__main__':
n, m = [int(i) for i in input().split()]
arr = [int(i) for i in input().split()]
STree = SegmentTree(n, arr)
STree.build()
for i in range(m):
t = [int(i) for i in input().split()]
if t[0] == 1:
STree.set(t[1], t[2])
else:
print(STree.first_above(t[1], t[2]))