n.py

New Zealand Programming Contest 2022 - Max Empty.py

# https://youtu.be/zo_UHPI78H8
# @robertkingNZ
# thred.co.nz
"""
Problem statement screenshot in comments below.
Problem N (page 29) https://prog4fun.csse.canterbury.ac.nz/pluginfile.php/11232/mod_quiz/intro/NZPC2022-v1.pdf

There's a clone of the contest on our prog4fun server: https://prog4fun.csse.canterbury.ac.nz.
This is a public server that anyone can register on with any credentials.
Once registered they should enrol themselves in the "course" Programming Contest Problem Archive and they'll then be able to access the clone of NZPC-2022 here:
https://prog4fun.csse.canterbury.ac.nz/mod/quiz/view.php?id=2796
"""

def hash(i, j, c):
    return i * c + j


def graph_from_grid(grid, r, c):
    n = r * c
    graph = [[] for _ in range(n)]

    def add(i, j, i2, j2):
        if grid[i][j] >= grid[i2][j2]:
            a = hash(i, j, c)
            b = hash(i2, j2, c)
            graph[a].append(b)

    for i in range(r):
        for j in range(1, c):
            add(i, j, i, j - 1)
        for j in range(c - 1):
            add(i, j, i, j + 1)
    for j in range(c):
        for i in range(1, r):
            add(i, j, i - 1, j)
        for i in range(r - 1):
            add(i, j, i + 1, j)
    height = [0] * n
    edges = set()
    for i in range(r):
        for j in range(c):
            a = hash(i, j, c)
            height[a] = grid[i][j]
            if i == 0 or j == 0 or i + 1 == r or j + 1 == c:
                edges.add(a)
    return graph, height, edges


class Dsu:
    def __init__(self, n):
        self.parent = list(range(n))
        self.size = [1] * n

    def find(self, i):
        p = self.parent[i]
        if p == i:
            return p
        self.parent[i] = self.find(p)
        return self.parent[i]

    def union(self, i, j):
        pi = self.find(i)
        pj = self.find(j)
        if pi == pj:
            return
        self.parent[pj] = pi
        self.size[pi] += self.size[pj]


def rec(i, height, area, tree, ph, pv, edges):
    my_area = area[i]
    dh = ph - height[i]
    my_volume = pv + my_area * dh
    ans = 0
    if i in edges:
        ans = my_volume
    for child in tree[i]:
        ans = max(ans, rec(child, height, area, tree, height[i], my_volume, edges))
    return ans


def solve():
    import sys
    sys.setrecursionlimit(500 * 500)
    data = (list(map(int, line.split())) for line in sys.stdin.readlines())
    c, r, h = next(data)
    c -= 2
    r -= 2
    grid = [next(data) for _ in range(r)]
    graph, height, edges = graph_from_grid(grid, r, c)
    n = len(height)
    dsu = Dsu(n)
    ord = list(range(n))
    ord.sort(key=height.__getitem__)
    visited = [False] * n
    tree = [[] for _ in range(n)]
    parents = set(range(n))
    for idx in ord:
        visited[idx] = True
        p = dsu.find(idx)
        for small_idx in graph[idx]:
            if visited[small_idx]:
                p2 = dsu.find(small_idx)
                if p2 != p:
                    tree[p].append(p2)
                    dsu.union(p, p2)
                    parents.remove(p2)
    assert (len(parents) == 1)
    i = next(iter(parents))
    ans = rec(i, height, dsu.size, tree, h - 1, 0, edges)
    print(ans)


solve()

here's the problem statement