Advertisement
Guest User

Untitled

a guest
Dec 11th, 2020
320
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 1.94 KB | None | 0 0
  1. import sys
  2. from typing import Dict, Set, List, Tuple
  3.  
  4. def parse_seats(lines: List[str]) -> Dict[Tuple[int, int], Set[Tuple[int, int]]]:
  5.     # find all seats in the input, and their connections.
  6.     w, h = len(lines[0]), len(lines)
  7.     seats = {(x, y): {(x, y)} for x in range(w) for y in range(h) if lines[y][x] != '.'}
  8.     for path in generate_paths(w, h):
  9.         previous = None
  10.         for current in path:
  11.             if current in seats:
  12.                 if previous:
  13.                     seats[previous].add(current)
  14.                     seats[current].add(previous)
  15.                 previous = current
  16.     return seats
  17.  
  18. def generate_paths(w: int, h: int):
  19.     # to the right
  20.     for y in range(0, h):
  21.         yield ((x, y) for x in range(w))
  22.     # down
  23.     for x in range(0, w):
  24.         yield ((x, y) for y in range(h))
  25.     # down and to the right
  26.     for x0 in range(-h + 1, w):
  27.         yield ((x0 + y, y) for y in range(max(0, -x0), min(h, w - x0)))
  28.     # down and to the left
  29.     for x0 in range(0, w + h):
  30.         yield ((x0 - y, y) for y in range(max(0, x0 - w + 1), min(h, x0 + 1)))
  31.  
  32. def solve(seats: Dict[Tuple[int, int], Set[Tuple[int, int]]], n: int):
  33.     # assume that all seats are free at the beginning.
  34.     free = {seat for seat, visible in seats.items() if len(visible) > n} # after 2 iterations
  35.     while True:
  36.         # visible = {seat, *neighbours}
  37.         changed = set(seat for seat, visible in seats.items() if visible <= free or seat not in free and len(visible - free) > n)
  38.         if not changed:
  39.             break
  40.         free ^= changed
  41.     return len(seats) - len(free)
  42.  
  43. if __name__ == '__main__':
  44.     lines = [line.rstrip() for line in sys.stdin]
  45.     seats_part2 = parse_seats(lines)
  46.     seats_part1 = {(i, j): {(p, q) for (p, q) in visible if max(abs(i-p), abs(j-q)) < 2} for (i, j), visible in seats_part2.items()}
  47.     print(f'Part 1: {solve(seats_part1, 4)}')
  48.     print(f'Part 2: {solve(seats_part2, 5)}')
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement