大和証券プログラミングコンテスト2024（AtCoder Beginner Contest 383）E,F,G問題メモ

from operator import itemgetter


class UnionFind:
    def __init__(self, n):
        self.table = [-1] * n
        self.group_a = []
        self.group_b = []

    def root(self, x):
        stack = []
        tbl = self.table
        while tbl[x] >= 0:
            stack.append(x)
            x = tbl[x]
        for y in stack:
            tbl[y] = x
        return x

    def find(self, x, y):
        return self.root(x) == self.root(y)

    def unite(self, x, y):
        r1 = self.root(x)
        r2 = self.root(y)
        if r1 == r2:
            return False
        d1 = self.table[r1]
        d2 = self.table[r2]
        if d1 <= d2:
            self.table[r2] = r1
            self.table[r1] += d2
            self.group_a[r1].update(self.group_a[r2])
            self.group_b[r1].update(self.group_b[r2])
        else:
            self.table[r1] = r2
            self.table[r2] += d1
            self.group_a[r2].update(self.group_a[r1])
            self.group_b[r2].update(self.group_b[r1])
        return True

    def get_size(self, x):
        return -self.table[self.root(x)]


def solve(n, m, k, edges, aaa, bbb):
    aaa = [a - 1 for a in aaa]
    bbb = [b - 1 for b in bbb]
    cnt_a = [0] * n
    cnt_b = [0] * n
    for a in aaa:
        cnt_a[a] += 1
    for b in bbb:
        cnt_b[b] += 1
    uft = UnionFind(n)
    uft.group_a = [{i} if cnt_a[i] > 0 else set() for i in range(n)]
    uft.group_b = [{i} if cnt_b[i] > 0 else set() for i in range(n)]
    edges.sort(key=itemgetter(2))
    ans = 0
    for u, v, w in edges:
        u -= 1
        v -= 1
        if uft.unite(u, v) == False:
            continue
        r = uft.root(u)
        ga = uft.group_a[r]
        gb = uft.group_b[r]
        while len(ga) > 0 and len(gb) > 0:
            a = next(iter(ga))
            b = next(iter(gb))
            ans += w
            cnt_a[a] -= 1
            if cnt_a[a] == 0:
                ga.remove(a)
            cnt_b[b] -= 1
            if cnt_b[b] == 0:
                gb.remove(b)
    return ans


n, m, k = map(int, input().split())
edges = [tuple(map(int, input().split())) for _ in range(m)]
aaa = list(map(int, input().split()))
bbb = list(map(int, input().split()))
ans = solve(n, m, k, edges, aaa, bbb)
print(ans)

from collections import defaultdict


def solve(n, x, k, items):
    items2 = defaultdict(list)
    for p, u, c in items:
        items2[c].append((p, u))
    dp = [0] * (x + 1)
    for itm in items2.values():
        ndp = dp.copy()
        for p, u in itm:
            for i in range(x, p - 1, -1):
                ndp[i] = max(ndp[i], ndp[i - p] + u, dp[i - p] + u + k)
        dp = ndp
    return dp[x]


n, x, k = map(int, input().split())
items = [tuple(map(int, input().split())) for _ in range(n)]
ans = solve(n, x, k, items)
print(ans)

import numpy as np
from numba import njit


@njit('(i8,i8,i8,i8[:])', cache=True)
def recursive(l, r, k, bbb):
    INF = 1 << 55
    w = r - l
    if w <= 2 * k:
        max_tile = (w - 1) // k + 1
        res = np.full((k, k, max_tile + 1), -INF, dtype=np.int64)
        res[:, :, 0] = 0
        for i in range(w):
            max_x = min(i, k - 1)
            max_y = min(w - i - 1, k - 1)
            for x in range(max_x + 1):
                for y in range(max_y + 1):
                    res[x, y, 1] = max(res[x, y, 1], bbb[l + i])
            for j in range(i + k, w):
                max_y = min(w - j - 1, k - 1)
                for x in range(max_x + 1):
                    for y in range(max_y + 1):
                        res[x, y, 2] = max(res[x, y, 2], bbb[l + i] + bbb[l + j])
        return res

    def merge(res1, res2, res, x1, y1, x2, y2):
        i1 = 1
        i2 = 1
        d1 = res1[x1, y1, 1]
        d2 = res2[x2, y2, 1]
        tmp = 0
        for i in range(1, res.shape[2]):
            if d1 >= d2:
                tmp += d1
                res[x1, y2, i] = max(res[x1, y2, i], tmp)
                i1 += 1
                if res1.shape[2] > i1:
                    d1 = res1[x1, y1, i1] - res1[x1, y1, i1 - 1]
                else:
                    d1 = -INF
            else:
                tmp += d2
                res[x1, y2, i] = max(res[x1, y2, i], tmp)
                i2 += 1
                if res2.shape[2] > i2:
                    d2 = res2[x2, y2, i2] - res2[x2, y2, i2 - 1]
                else:
                    d2 = -INF

    m = (l + r) // 2
    res1 = recursive(l, m, k, bbb)
    res2 = recursive(m, r, k, bbb)
    max_tile = (r - l - 1) // k + 1
    res = np.full((k, k, max_tile + 1), -INF, dtype=np.int64)
    res[:, :, 0] = 0
    for x in range(k):
        for y in range(k):
            for z in range(k):
                merge(res1, res2, res, x, z, k - z - 1, y)
    for t in range(1, max_tile + 1):
        for x in range(k):
            for y in range(k - 1, 0, -1):
                res[x, y - 1, t] = max(res[x, y - 1, t], res[x, y, t])
        for y in range(k):
            for x in range(k - 1, 0, -1):
                res[x - 1, y, t] = max(res[x - 1, y, t], res[x, y, t])
    return res


@njit('(i8,i8,i8[:])', cache=True)
def solve(n, k, aaa):
    bbb = np.zeros(n - k + 1, dtype=np.int64)
    s = sum(aaa[:k])
    bbb[0] = s
    for i in range(k, n):
        s += aaa[i] - aaa[i - k]
        bbb[i - k + 1] = s
    res = recursive(0, n - k + 1, k, bbb)

    ans = np.zeros(n // k, np.int64)
    for i in range(1, n // k + 1):
        ans[i - 1] = res[:, :, i].max()
    return ans


n, k = map(int, input().split())
aaa = np.fromstring(input(), dtype=np.int64, sep=' ')
ans = solve(n, k, aaa)
print(*ans)