Advent of Code 2020 - Day 11: Seating System

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Advent of Code 2020 - Day 11: Seating System
https://adventofcode.com/2020/day/11

:license: GPLv3 --- Copyright (C) 2020 Olivier Pirson
:author: Olivier Pirson --- http://www.opimedia.be/
:version: December 11, 2020
"""

import sys

from typing import List, Tuple


def nb_place(grid: List[str], x: int, y: int, place: str) -> int:
    height = len(grid)
    width = len(grid[0])

    assert 0 <= x < width
    assert 0 <= y < height
    assert place in '.#L'

    nb = 0

    for y_offset in (-1, 0, 1):
        y_new = y + y_offset
        if 0 <= y_new < height:
            row = grid[y_new]
            for x_offset in (-1, 0, 1):
                if (x_offset != 0) or (y_offset != 0):
                    x_new = x + x_offset
                    if (0 <= x_new < width) and (row[x_new] == place):
                        nb += 1

    return nb


def nb_visible(grid: List[str], x: int, y: int, place: str) -> int:
    height = len(grid)
    width = len(grid[0])

    assert 0 <= x < width
    assert 0 <= y < height
    assert place in '.#L'

    nb = 0

    for y_offset in (-1, 0, 1):
        for x_offset in (-1, 0, 1):
            if (x_offset != 0) or (y_offset != 0):
                x_new = x + x_offset
                y_new = y + y_offset
                while ((0 <= x_new < width) and (0 <= y_new < height)
                       and (grid[y_new][x_new] == '.')):
                    x_new += x_offset
                    y_new += y_offset
                if ((0 <= x_new < width) and (0 <= y_new < height)
                        and (grid[y_new][x_new] == place)):
                    nb += 1

    return nb


def update1(grid: List[str]) -> Tuple[List[str], bool]:
    new_grid = []
    is_modified = False

    for y, row in enumerate(grid):
        new_row = []
        for x, place in enumerate(row):
            if place == 'L':
                nb = nb_place(grid, x, y, '#')
                if nb == 0:
                    is_modified = True
                    place = '#'
            elif place == '#':
                nb = nb_place(grid, x, y, '#')
                if nb >= 4:
                    is_modified = True
                    place = 'L'
            new_row.append(place)

        new_grid.append(''.join(new_row))

    return new_grid, is_modified


def update2(grid: List[str]) -> Tuple[List[str], bool]:
    new_grid = []
    is_modified = False

    for y, row in enumerate(grid):
        new_row = []
        for x, place in enumerate(row):
            if place == 'L':
                nb = nb_visible(grid, x, y, '#')
                if nb == 0:
                    is_modified = True
                    place = '#'
            elif place == '#':
                nb = nb_visible(grid, x, y, '#')
                if nb >= 5:
                    is_modified = True
                    place = 'L'
            else:
                nb = 0
            new_row.append(place)

        new_grid.append(''.join(new_row))

    return new_grid, is_modified


def main() -> None:
    grid = []
    for line in sys.stdin:
        grid.append(line.rstrip())

    # Part 1
    new_grid = grid

    is_modified = True
    while is_modified:
        is_modified = False
        new_grid, is_modified = update1(new_grid)

    nb = sum(row.count('#') for row in new_grid)
    print(nb)

    # Part 2
    new_grid = grid

    is_modified = True
    while is_modified:
        is_modified = False
        new_grid, is_modified = update2(new_grid)
        nb = sum(row.count('#') for row in new_grid)

    nb = sum(row.count('#') for row in new_grid)
    print(nb)


if __name__ == '__main__':
    main()