AtCoder Beginner Contest 347 F,G問題メモ

❷を固定すれば左上と右下からだけでいい、ということに気付かず、4隅から累積MAXを取っている。

def accumulate2d(aaa):
    n = len(aaa)
    acc = [row.copy() for row in aaa]
    for i in range(n):
        for j in range(1, n):
            acc[i][j] = max(acc[i][j], acc[i][j - 1])
    for j in range(n):
        for i in range(1, n):
            acc[i][j] = max(acc[i][j], acc[i - 1][j])
    return acc


def rot90r(aaa):
    return list(map(list, zip(*reversed(aaa))))


def rot90l(aaa):
    return list(reversed(list(map(list, zip(*aaa)))))


def rot180(aaa):
    return [list(reversed(row)) for row in reversed(aaa)]


n, m = map(int, input().split())
grid = [list(map(int, input().split())) for _ in range(n)]
acc = [[0] * (n + 1)]
for row in grid:
    acc.append([0] + row)
for i in range(1, n + 1):
    for j in range(1, n + 1):
        acc[i][j] += acc[i][j - 1]
for j in range(1, n + 1):
    for i in range(1, n + 1):
        acc[i][j] += acc[i - 1][j]

squares = [[0] * (n - m + 1) for _ in range(n - m + 1)]
for i in range(m, n + 1):
    for j in range(m, n + 1):
        squares[i - m][j - m] = acc[i][j] - acc[i][j - m] - acc[i - m][j] + acc[i - m][j - m]

acc_lu = accumulate2d(squares)
acc_ld = rot90l(accumulate2d(rot90r(squares)))
acc_ru = rot90r(accumulate2d(rot90l(squares)))
acc_rd = rot180(accumulate2d(rot180(squares)))
acc_row = list(map(max, squares))
acc_col = list(map(max, zip(*squares)))

ans = 0
nm1 = n - m + 1
for i in range(m, nm1 - m):
    ans = max(ans,
              acc_row[i] + acc_lu[i - m][-1] + acc_ld[i + m][-1],
              acc_col[i] + acc_lu[-1][i - m] + acc_ru[-1][i + m])

for i in range(nm1 - m):
    for j in range(nm1 - m):
        ans = max(ans,
                  acc_lu[i][j] + acc_ru[i][j + m] + acc_ld[i + m][-1],
                  acc_lu[i][j] + acc_ld[i + m][j] + acc_ru[-1][j + m],
                  acc_ru[i][j + m] + acc_rd[i + m][j + m] + acc_lu[-1][j],
                  acc_lu[i][-1] + acc_ld[i + m][j] + acc_rd[i + m][j + m], )
print(ans)

from collections import deque
from typing import Tuple, List


class Dinic:
    """
    Usage:
       mf = Dinic(n)
    -> mf.add_link(from, to, capacity)
    -> mf.max_flow(source, target)
    """

    def __init__(self, n: int):
        self.n = n
        self.links: List[List[List[int]]] = [[] for _ in range(n)]
        # if exists an edge (v→u, capacity)...
        #   links[v] = [ [ capacity, u, index of rev-edge in links[u], is_original_edge ], ]

    def add_link(self, from_: int, to: int, capacity: int) -> None:
        self.links[from_].append([capacity, to, len(self.links[to]), 1])
        self.links[to].append([0, from_, len(self.links[from_]) - 1, 0])

    def bfs(self, s: int) -> List[int]:
        depth = [-1] * self.n
        depth[s] = 0
        q = deque([s])
        while q:
            v = q.popleft()
            for cap, to, rev, _ in self.links[v]:
                if cap > 0 and depth[to] < 0:
                    depth[to] = depth[v] + 1
                    q.append(to)
        return depth

    def dfs(self, s: int, t: int, depth: List[int], progress: List[int], link_counts: List[int]) -> int:
        links = self.links
        stack = [s]

        while stack:
            v = stack[-1]
            if v == t:
                break
            for i in range(progress[v], link_counts[v]):
                progress[v] = i
                cap, to, rev, _ = links[v][i]
                if cap == 0 or depth[v] >= depth[to] or progress[to] >= link_counts[to]:
                    continue
                stack.append(to)
                break
            else:
                progress[v] += 1
                stack.pop()
        else:
            return 0

        f = 1 << 60
        fwd_links = []
        bwd_links = []
        for v in stack[:-1]:
            cap, to, rev, _ = link = links[v][progress[v]]
            f = min(f, cap)
            fwd_links.append(link)
            bwd_links.append(links[to][rev])

        for link in fwd_links:
            link[0] -= f

        for link in bwd_links:
            link[0] += f

        return f

    def max_flow(self, s: int, t: int) -> int:
        link_counts = list(map(len, self.links))
        flow = 0
        while True:
            depth = self.bfs(s)
            if depth[t] < 0:
                break
            progress = [0] * self.n
            current_flow = self.dfs(s, t, depth, progress, link_counts)
            while current_flow > 0:
                flow += current_flow
                current_flow = self.dfs(s, t, depth, progress, link_counts)
        return flow

    def cut_edges(self, s: int) -> List[Tuple[int, int]]:
        """ max_flowしたあと、最小カットにおいてカットすべき辺を復元する """
        q = [s]
        reachable = [0] * self.n
        reachable[s] = 1
        while q:
            v = q.pop()
            for cap, u, li, _ in self.links[v]:
                if cap == 0 or reachable[u]:
                    continue
                reachable[u] = 1
                q.append(u)
        edges = []
        for v in range(self.n):
            if reachable[v] == 0:
                continue
            for cap, u, li, orig in self.links[v]:
                if orig == 1 and reachable[u] == 0:
                    edges.append((v, u))
        return edges


n = int(input())
aaa = [list(map(int, input().split())) for _ in range(n)]

s = n * n * 4
t = s + 1
dnc = Dinic(t + 1)
INF = 1 << 60


def set_links(dnc: Dinic, v0: int, v1: int):
    dnc.add_link(v0 + 0, v1 + 0, 1)
    dnc.add_link(v0 + 1, v1 + 1, 1)
    dnc.add_link(v0 + 2, v1 + 2, 1)
    dnc.add_link(v0 + 3, v1 + 3, 1)
    dnc.add_link(v0 + 1, v1 + 0, 2)
    dnc.add_link(v0 + 2, v1 + 1, 2)
    dnc.add_link(v0 + 3, v1 + 2, 2)
    dnc.add_link(v0 + 2, v1 + 0, 2)
    dnc.add_link(v0 + 3, v1 + 1, 2)
    dnc.add_link(v0 + 3, v1 + 0, 2)


for i in range(n):
    for j in range(n):
        a = aaa[i][j]
        v0 = (i * n + j) * 4
        if a > 0:
            if a != 1:
                dnc.add_link(s, v0 + 0, INF)
            if a != 2:
                dnc.add_link(v0 + 0, v0 + 1, INF)
            if a != 3:
                dnc.add_link(v0 + 1, v0 + 2, INF)
            if a != 4:
                dnc.add_link(v0 + 2, v0 + 3, INF)
            if a != 5:
                dnc.add_link(v0 + 3, t, INF)
        dnc.add_link(v0 + 3, v0 + 2, INF)
        dnc.add_link(v0 + 2, v0 + 1, INF)
        dnc.add_link(v0 + 1, v0 + 0, INF)
        if i > 0:
            v1 = ((i - 1) * n + j) * 4
            set_links(dnc, v0, v1)
        if i < n - 1:
            v1 = ((i + 1) * n + j) * 4
            set_links(dnc, v0, v1)
        if j > 0:
            v1 = (i * n + j - 1) * 4
            set_links(dnc, v0, v1)
        if j < n - 1:
            v1 = (i * n + j + 1) * 4
            set_links(dnc, v0, v1)

cost = dnc.max_flow(s, t)
edges = dnc.cut_edges(s)

ans1 = [[6] * n for _ in range(n)]
ans2 = [[0] * n for _ in range(n)]

for u, v in edges:
    uij, uk = divmod(u, 4)
    vij, vk = divmod(v, 4)
    ui, uj = divmod(uij, n)
    vi, vj = divmod(vij, n)
    ans1[vi][vj] = min(ans1[vi][vj], vk + 1)
    ans2[ui][uj] = max(ans2[ui][uj], uk + 2)

for i in range(n):
    for j in range(n):
        if aaa[i][j] != 0:
            continue
        if ans1[i][j] == 6:
            aaa[i][j] = ans2[i][j]
        else:
            aaa[i][j] = ans1[i][j]


def dfs(si, sj, move=((-1, 0), (1, 0), (0, -1), (0, 1))):
    a = aaa[si][sj]
    q = [(si, sj)]
    while q:
        i, j = q.pop()
        for di, dj in move:
            ni = i + di
            nj = j + dj
            if 0 <= ni < n and 0 <= nj < n:
                if aaa[ni][nj] == 0:
                    aaa[ni][nj] = a
                    q.append((ni, nj))


for i in range(n):
    for j in range(n):
        if aaa[i][j] != 0:
            dfs(i, j)

for row in aaa:
    print(*row)