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

どのような選び方 $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)))

from atcoder.dsu import DSU


def solve(f, h, w, grids, q, queries):
    uft = DSU(f * h * w)
    for k in range(f):
        offset = k * h * w
        for i in range(h):
            for j in range(w):
                if grids[k][i][j] == '#':
                    continue
                if i < h - 1 and grids[k][i + 1][j] != '#':
                    uft.merge(offset + i * w + j, offset + (i + 1) * w + j)
                if j < w - 1 and grids[k][i][j + 1] != '#':
                    uft.merge(offset + i * w + j, offset + i * w + (j + 1))

    leaders = []
    for k in range(f):
        # dbg_floor_leaders = []
        for i in range(h):
            for j in range(w):
                if grids[k][i][j] == '#':
                    continue
                key = k * h * w + i * w + j
                if uft.leader(key) == key:
                    leaders.append(key)
                    # dbg_floor_leaders.append((k, i, j))
        # print(f'{k=} {dbg_floor_leaders=}')

    n = len(leaders)
    mapping = {v: i for i, v in enumerate(leaders)}

    floor_scores = [0] * n
    links = [set() for _ in range(n)]
    rev_links = [set() for _ in range(n)]
    for k in range(f):
        offset1 = k * h * w
        offset2 = (k + 1) * h * w
        for i in range(h):
            for j in range(w):
                if grids[k][i][j] == '#':
                    continue
                s = int(grids[k][i][j])
                r = mapping[uft.leader(offset1 + i * w + j)]
                floor_scores[r] += s
                if k < f - 1 and grids[k + 1][i][j] != '#':
                    r2 = mapping[uft.leader(offset2 + i * w + j)]
                    links[r].add(r2)
                    rev_links[r2].add(r)

    # はしごを使わず、そこからの移動で取得できる (最大値, 最大値を与える次の移動先, 2番目の値)
    scores_no_return = [(0, -1, 0)] * n
    for u in range(n - 1, -1, -1):
        fsu = floor_scores[u]
        if len(links[u]) == 0:
            scores_no_return[u] = (fsu, -1, 0)
            continue
        if len(links[u]) == 1:
            v = next(iter(links[u]))
            scores_no_return[u] = (fsu + scores_no_return[v][0], v, 0)
            continue
        v, x = sorted(links[u], key=lambda v: scores_no_return[v][0], reverse=True)[:2]
        scores_no_return[u] = (fsu + scores_no_return[v][0], v, fsu + scores_no_return[x][0])

    # その場ではしごを使ってそこからの移動で取得できる (最大値, 最大値を与えるはしごでの移動先, 2番目の値)
    scores_right_away = [(0, -1, 0)] * n
    for v in range(n):
        if len(rev_links[v]) == 0:
            continue
        fsv = floor_scores[v]
        best_score = 0
        best_u = -1
        second_score = 0
        for u in rev_links[v]:
            su1, nxt, su2 = scores_no_return[u]
            if nxt != v:
                tmp = fsv + su1
            else:
                tmp = max(su1, fsv + su2)
            if tmp > best_score:
                second_score = best_score
                best_score = tmp
                best_u = u
            elif tmp > second_score:
                second_score = tmp
        scores_right_away[v] = (best_score, best_u, second_score)

    # その場でははしごを使わず、その後に高々1回使って、そこからの移動で取得できる最大値
    scores_later = [0] * n
    for u in range(n - 1, -1, -1):
        fsu = floor_scores[u]
        best = fsu
        for v in links[u]:
            sr1, nxt, sr2 = scores_right_away[v]
            slt = scores_later[v]
            if nxt != u:
                tmp = fsu + max(sr1, slt)
            else:
                tmp = max(sr1, fsu + sr2, fsu + slt)
            best = max(best, tmp)
        scores_later[u] = best

    buf = []
    for g, a, b in queries:
        g -= 1
        a -= 1
        b -= 1
        v = mapping[uft.leader(g * h * w + a * w + b)]
        ans = max(
            scores_right_away[v][0],
            scores_later[v],
        )
        buf.append(ans)
    return buf


f, h, w = list(map(int, input().split()))
grids = []
for _ in range(f):
    grids.append([input() for _ in range(h)])
q = int(input())
queries = [list(map(int, input().split())) for _ in range(q)]
ans = solve(f, h, w, grids, q, queries)
print('\n'.join(map(str, ans)))