USACO Castle

castle.py

'''
ID: daniel20
LANG: PYTHON3
TASK: castle
'''

global WEST, NORTH, EAST, SOUTH
WEST, NORTH, EAST, SOUTH = 0, 1, 2, 3

def closed(room, direction=NORTH):
    return bool(room & (1 << direction))
def opened(room, direction=NORTH):
    return not closed(room, direction)

def connected_components(graph):
    components = dict() # vertex label -> component label
    sizes = dict() # component label -> size of component

    def in_bounds(index, array):
        return 0 <= index < len(array)
    def dfs(graph, r, c, label):
        if (not in_bounds(r, graph)) or (not in_bounds(c, graph[0])):
            return
        if (r, c) in components:
            return
        
        components[(r, c)] = label
        sizes[label] += 1

        room = graph[r][c]
        if opened(room, direction=NORTH):
            dfs(graph, r - 1, c, label)
        if opened(room, direction=SOUTH):
            dfs(graph, r + 1, c, label)
        if opened(room, direction=WEST):
            dfs(graph, r, c - 1, label)
        if opened(room, direction=EAST):
            dfs(graph, r, c + 1, label)
    
    num_components = 0
    for r in range(len(graph)):
        for c in range(len(graph[0])):
            if (r, c) not in components:
                num_components += 1

                sizes[num_components] = 0
                dfs(graph, r, c, num_components)
    
    return components, sizes

def main():
    fin = open('castle.in', 'r')
    fout = open('castle.out', 'w')
    
    # read the input
    M, N = map(int, fin.readline().split(' '))
    castle = list()
    for i in range(N):
        row = list(map(int, fin.readline().split(' ')))
        castle.append(row)
    fin.close()

    components, sizes = connected_components(castle)
    
    num_rooms = len(sizes)
    largest_room_size = sizes[max(sizes, key=sizes.get)]

    # iterate through all possible edges, try adding each one (~ V(V-1)/2 edges)
    # a valid edge must be between two adjacent vertices and end vertices must be in diff components

    def adjacent(u, v):
        ur, uc = u
        vr, vc = v

        rdiff, cdiff = abs(ur - vr), abs(uc - vc)
        if rdiff == 1 and cdiff == 0:
            return True
        if rdiff == 0 and cdiff == 1:
            return True
        return False

    # max component size by adding an edge
    largest_room_size_removal = -1
    wall = None

    for ur in range(len(castle)):
        for uc in range(len(castle[0])):
            v = [(ur + 1, uc), (ur - 1, uc), (ur, uc + 1), (ur, uc - 1)]
            for vr, vc in v:
                # node 1, node 2
                n1, n2 = (ur, uc), (vr, vc)
                if n1 == n2:
                    continue
                if vr < 0 or vr >= len(castle) or vc < 0 or vc >= len(castle[0]):
                    continue
                if not adjacent(n1, n2):
                    continue
                if components[n1] == components[n2]:
                    continue

                # current component size (combined)
                current_component_size = sizes[components[n1]] + sizes[components[n2]]
                
                # find the wall coordinates (and direction)
                wall_r, wall_c = 0, 0
                wall_dir = NORTH
                if ur == vr:
                    if uc < vc:
                        wall_r, wall_c = ur + 1, uc + 1
                    else:
                        wall_r, wall_c = vr + 1, vc + 1
                    wall_dir = EAST
                else:
                    if ur > vr:
                        wall_r, wall_c = ur + 1, uc + 1
                    else:
                        wall_r, wall_c = vr + 1, vc + 1
                    wall_dir = NORTH

                if current_component_size == largest_room_size_removal:
                    r, c, dir = wall
                    if wall_c < c:
                        wall = (wall_r, wall_c, wall_dir)
                    elif c < wall_c:
                        pass
                    else:
                        if wall_r > r:
                            wall = (wall_r, wall_c, wall_dir)
                        elif r > wall_r:
                            pass
                        else:
                            if wall_dir == NORTH:
                                wall = (wall_r, wall_c, wall_dir)
                            else:
                                pass
                    
                    r, c, dir = wall
                else:
                    if current_component_size > largest_room_size_removal:
                        largest_room_size_removal = current_component_size
                        wall = (wall_r, wall_c, wall_dir)

    def println(val):
        fout.write(f'{val}\n')
    
    println(num_rooms)
    println(largest_room_size)
    println(largest_room_size_removal)

    dir_to_str = lambda dir: 'N' if dir == NORTH else ('S' if dir == SOUTH else ('W' if dir == WEST else 'E'))
    println(f'{wall[0]} {wall[1]} {dir_to_str(wall[2])}')
    
    fout.close()

if __name__ == '__main__':
    main()