USACO Dec 2015 Silver Problem #1 solution for CompSci club

lightson.py

#!/usr/bin/env python3

from collections import defaultdict

__author__ = 'Yatharth Agarwal <[email protected]>'

with open('lightson.in') as in_file:
    read_numbers = lambda: map(int, in_file.readline().strip().split())

    switches = defaultdict(set)
    N, M = read_numbers()

    for _ in range(M):
        x, y, a, b = read_numbers()
        switches[(x, y)].add((a, b))

START_ROOM = (1, 1)

queue = {START_ROOM}
visited = set()
lit = {START_ROOM}


def neighbours(room):
    x, y = room
    for neighbour in {(x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)}:
        neighbour_x, neighbour_y = neighbour
        if 1 <= neighbour_x <= N and 1 <= neighbour_y <= N:
            yield neighbour


def hasVisitedNeighbour(room):
    return any(neighbour in visited for neighbour in neighbours(room))


while queue:
    room = queue.pop()

    if room in visited:
        continue

    visited.add(room)

    for lit_room in switches[room]:
        lit.add(lit_room)

        if hasVisitedNeighbour(lit_room):
            queue.add(lit_room)

    for neighbour in neighbours(room):
        if neighbour in lit:
            queue.add(neighbour)

with open('lightson.out', 'w') as out_file:
    out_file.write("{}\n".format(len(lit)))

lightson.txt

Farmer John has recently built an enormous barn consisting of an N×N grid of rooms (2≤N≤100), numbered from (1,1) up to (N,N). Being somewhat afraid of the dark, Bessie the cow wants to turn on the lights in as many rooms as possible.
Bessie starts in room (1,1), the only room that is initially lit. In some rooms, she will find light switches that she can use to toggle the lights in other rooms; for example there might be a switch in room (1,1) that toggles the lights in room (1,2). Bessie can only travel through lit rooms, and she can only move from a room (x,y) to its four adjacent neighbors (x−1,y), (x+1,y), (x,y−1) and (x,y+1) (or possibly fewer neighbors if this room is on the boundary of the grid).

Please determine the maximum number of rooms Bessie can illuminate.

INPUT FORMAT (file lightson.in):

The first line of input contains integers N and M (1≤M≤20,000).
The next MM lines each describe a single light switch with four integers x, y, a, b, that a switch in room (x,y) can be used to toggle the lights in room (a,b). Multiple switches may exist in any room, and multiple switches may toggle the lights of any room.

OUTPUT FORMAT (file lightson.out):

A single line giving the maximum number of rooms Bessie can illuminate.
SAMPLE INPUT:

 
3 6
1 1 1 2
2 1 2 2
1 1 1 3
2 3 3 1
1 3 1 2
1 3 2 1

SAMPLE OUTPUT:
5

Here, Bessie can use the switch in (1,1) to turn on lights in (1,2) and (1,3). She can then walk to (1,3) and turn on the lights in (2,1), from which she can turn on the lights in (2,2). The switch in (2,3) is inaccessible to her, being in an unlit room. She can therefore illuminate at most 5 rooms.

Problem credits: Austin Bannister and Brian Dean