python code_festival_2018_final_f

Posted 2021-05-09

import sys
from heapq import heappush, heappushpop, heappop


class SegmentTree:
    def __init__(self, n):
        d = (n - 1).bit_length()
        self.n = 1 << d
        self.offset = self.n - 1
        self.segMin = [[0, 0] for _ in range(self.n << 1)]
        self.segMax = [[0, 0] for _ in range(self.n << 1)]
        self.segAdd = [0] * (self.n << 1)
        for i in range(self.n):
            self.segMin[i + self.offset][1] = i
            self.segMax[i + self.offset][1] = i

    def add(self, a, b, x):
        return self._add(a, b, x, 0, 0, self.n)

    def _add(self, a, b, x, k, l, r):
        if r <= a or b <= l:
            return

        if a <= l and r <= b:
            self.segAdd[k] += x
            return

        nk = k << 1
        m = (l + r) // 2
        self._add(a, b, x, nk + 1, l, m)
        self._add(a, b, x, nk + 2, m, r)

        la = self.segAdd[nk + 1]
        ra = self.segAdd[nk + 2]
        lv, li = self.segMin[nk + 1]
        rv, ri = self.segMin[nk + 2]
        mk = self.segMin[k]
        if lv + la < rv + ra:
            mk[0] = lv + la
            mk[1] = li
        else:
            mk[0] = rv + ra
            mk[1] = ri
        lv, li = self.segMax[nk + 1]
        rv, ri = self.segMax[nk + 2]
        mk = self.segMax[k]
        if lv + la > rv + ra:
            mk[0] = lv + la
            mk[1] = li
        else:
            mk[0] = rv + ra
            mk[1] = ri

    def getMax(self, a, b):
        return self._get(a, b, 0, 0, self.n, self.segMax, float('-inf'), max, True)

    def getMin(self, a, b):
        return self._get(a, b, 0, 0, self.n, self.segMin, float('inf'), min, False)

    def _get(self, a, b, k, l, r, data, default, func, is_min):
        if r <= a or b <= l:
            return default, 0

        if a <= l and r <= b:
            # self.lazy(k, data, is_min)
            v, i = data[k]
            return v + self.segAdd[k], i

        nk = k << 1
        m = (l + r) // 2
        lvi = self._get(a, b, nk + 1, l, m, data, default, func, is_min)
        rvi = self._get(a, b, nk + 2, m, r, data, default, func, is_min)
        ret = func(lvi, rvi)

        return ret[0] + self.segAdd[k], ret[1]

    def lazy(self, k, data, is_min):
        print(k, data[k], self.segAdd[k])
        if k >= self.offset:
            return data[k]
        nk = k << 1
        lv, li = self.lazy(nk + 1, data, is_min)
        rv, ri = self.lazy(nk + 2, data, is_min)

        la = self.segAdd[nk + 1]
        ra = self.segAdd[nk + 2]
        mk = data[k]
        lv += la
        rv += ra
        if (is_min and lv < rv) or (not is_min and lv > rv):
            mk[0] = lv
            mk[1] = li
        else:
            mk[0] = rv
            mk[1] = ri
        return mk

    def debug_print(self):
        print(*(self.getMin(i, i + 1) for i in range(self.n)))


n = 15
st = SegmentTree(n + 1)
st.add(1, 5, 10)
st.add(4, 8, 5)
st.debug_print()
print(st.getMin(1, 8))
print(st.getMin(3, 6))
print(st.getMax(1, 8))
exit()

n, k = map(int, input().split())
ans = 0
ramen = []
restaurant = []
st = SegmentTree(n + 1)
frontier_0 = 0
frontier_k = 0
gourmet = 0

for i, line in enumerate(sys.stdin.readlines()):
    t, a = map(int, line.split())
    # st.debug_print()
    # print(i, ans, f't={t}', f'a={a}', ramen, restaurant)
    ans += a
    if t == 0:
        st.add(i + 1, st.n, -1)
        if gourmet > 0:
            heappush(ramen, (a, i))
            gourmet -= 1
        else:
            frontier_0 = i + 1
            b, j = heappushpop(ramen, (a, i))
            while j < frontier_k:
                b, j = heappop(ramen)
            ans -= b
            st.add(j + 1, st.n, 1)
            new_max, m = st.getMax(frontier_k, st.n)
            if new_max == k:
                frontier_k = m
    else:
        st.add(i + 1, st.n, 1)
        if gourmet < k:
            heappush(restaurant, (a, i))
            gourmet += 1
        else:
            frontier_k = i + 1
            b, j = heappushpop(restaurant, (a, i))
            while j < frontier_0:
                b, j = heappop(restaurant)
            ans -= b
            st.add(j + 1, st.n, -1)
            new_min, m = st.getMin(frontier_0, st.n)
            if new_min == 0:
                frontier_0 = m
print(ans)