AtCoder Beginner Contest 440 (Promotion of Engineer Guild Fes) E,F問題メモ

どのような選び方 $C_{cur}$ についても、それに遷移する1つ前の選び方 $C_{prev}$ が一意に定まることを示せばよい。
（始点 $C_{cur}=(K,0,...,0)$ の場合は除く）

この遷移のルールでは、以下が成り立つ。

• $C_{cur}$ の中で $0$ でない最も右のindexを $i$ とする
• $c_i$ を1減らし、$c_{i-1}$ を1増やしたものが、$C_{prev}$ である

親の辿り方が一意に定まるので、探索は重複していないことが示せる。

from heapq import heappop, heappush


def solve(n, k, x, aaa):
    if n == 1:
        return [aaa[0] * k]
    aaa.sort(reverse=True)
    ans = []
    q = [(-aaa[0] * k, 0, k, 0)]
    for _ in range(x):
        score, i, c1, c2 = heappop(q)
        score = -score
        ans.append(score)
        if c1 > 0:
            ns = score - aaa[i] + aaa[i + 1]
            heappush(q, (-ns, i, c1 - 1, c2 + 1))
        if i < n - 2 and c2 > 0:
            ns = score - aaa[i + 1] + aaa[i + 2]
            heappush(q, (-ns, i + 1, c2 - 1, 1))
    return ans


n, k, x = list(map(int, input().split()))
aaa = list(map(int, input().split()))
ans = solve(n, k, x, aaa)
print('\n'.join(map(str, ans)))

from operator import add, lt


class FenwickTreeInjectable:
    def __init__(self, n, identity_factory, func):
        self.size = n
        self.tree = [identity_factory() for _ in range(n + 1)]
        self.func = func
        self.idf = identity_factory
        self.depth = n.bit_length()

    def add(self, i, x):
        i += 1
        tree = self.tree
        func = self.func
        while i <= self.size:
            tree[i] = func(tree[i], x)
            i += i & -i

    def sum(self, i):
        i += 1
        s = self.idf()
        tree = self.tree
        func = self.func
        while i > 0:
            s = func(s, tree[i])
            i -= i & -i
        return s

    def lower_bound(self, x, lt):
        """
        累積和がx以上になる最小のindexと、その直前までの累積和

        :param lt: lt(a, b) で a < b ならTrueを返す関数
        """
        total = self.idf()
        pos = 0
        tree = self.tree
        func = self.func
        for i in range(self.depth, -1, -1):
            k = pos + (1 << i)
            if k > self.size:
                continue
            new_total = func(total, tree[k])
            if lt(new_total, x):
                total = new_total
                pos += 1 << i
        return pos, total

    def debug_print(self):
        prev = 0
        arr = []
        for i in range(self.size):
            curr = self.sum(i)
            arr.append(curr - prev)
            prev = curr
        print(arr)


from collections import Counter
from heapq import heapify, heappop, heappush
from typing import List


class RemovableHeapQueue:
    q: List[int]
    count: int
    total: int
    term_counter: Counter

    def __init__(self, arr: List[int] = None):
        if arr is None:
            self.q = []
            self.count = 0
            self.total = 0
            self.term_counter = Counter()
        else:
            self.q = arr
            heapify(self.q)
            self.count = len(arr)
            self.total = sum(arr)
            self.term_counter = Counter(arr)

    def push(self, x):
        heappush(self.q, x)
        self.count += 1
        self.total += x
        self.term_counter[x] += 1

    def pop(self):
        x = heappop(self.q)
        self.count -= 1
        self.total -= x
        self.term_counter[x] -= 1
        self.modify()
        return x

    def get_top(self):
        return self.q[0]

    def remove(self, x):
        assert self.term_counter[x] > 0
        self.count -= 1
        self.total -= x
        self.term_counter[x] -= 1
        self.modify()

    def modify(self):
        while self.q and self.term_counter[self.q[0]] == 0:
            heappop(self.q)

    def debug_print(self):
        print(f'cnt:{self.count} tot:{self.total} q:{self.q}')
        print(self.term_counter)


def solve(n, q, horses, queries):
    values = {0}
    for a, b in horses:
        values.add(-a)
    for w, x, y in queries:
        values.add(-x)
    values = sorted(values)
    mapping = {v: i for i, v in enumerate(values)}
    m = len(values)

    mood_all_count = FenwickTreeInjectable(m, int, add)
    mood_all_total = FenwickTreeInjectable(m, int, add)
    mood_neat = RemovableHeapQueue()
    for a, b in horses:
        ai = mapping[-a]
        mood_all_count.add(ai, 1)
        mood_all_total.add(ai, a)
        if b == 2:
            mood_neat.push(a)

    buf = []
    for w, x, y in queries:
        w -= 1
        pa, pb = horses[w]
        pi = mapping[-pa]
        mood_all_count.add(pi, -1)
        mood_all_total.add(pi, -pa)
        if pb == 2:
            mood_neat.remove(pa)
        xi = mapping[-x]
        mood_all_count.add(xi, 1)
        mood_all_total.add(xi, x)
        if y == 2:
            mood_neat.push(x)
        horses[w] = (x, y)

        k = min(n - 1, mood_neat.count)
        if k == 0:
            buf.append(mood_all_total.sum(m - 1))
            continue

        # 丁寧な馬の中での機嫌最小値は、一時的に削除
        mn = mood_neat.get_top()
        mni = mapping[-mn]
        mood_all_count.add(mni, -1)
        mood_all_total.add(mni, -mn)

        # 機嫌上位 k 個と、それ以外のそれぞれの和
        ki, l = mood_all_count.lower_bound(k, lt)
        kv = -values[ki]
        double = mood_all_total.sum(ki - 1) if ki > 0 else 0
        double += kv * (k - l)
        single = mood_all_total.sum(m - 1) - double + mn
        ans = single + double * 2
        buf.append(ans)

        # 削除したものを戻す
        mood_all_count.add(mni, 1)
        mood_all_total.add(mni, mn)

    return buf


n, q = list(map(int, input().split()))
horses = [list(map(int, input().split())) for _ in range(n)]
queries = [list(map(int, input().split())) for _ in range(q)]
ans = solve(n, q, horses, queries)
print('\n'.join(map(str, ans)))