Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import sys
- from typing import Dict, Set, List, Tuple
- def parse_seats(lines: List[str]) -> Dict[Tuple[int, int], Set[Tuple[int, int]]]:
- # find all seats in the input, and their connections.
- w, h = len(lines[0]), len(lines)
- seats = {(x, y): {(x, y)} for x in range(w) for y in range(h) if lines[y][x] != '.'}
- for path in generate_paths(w, h):
- previous = None
- for current in path:
- if current in seats:
- if previous:
- seats[previous].add(current)
- seats[current].add(previous)
- previous = current
- return seats
- def generate_paths(w: int, h: int):
- # to the right
- for y in range(0, h):
- yield ((x, y) for x in range(w))
- # down
- for x in range(0, w):
- yield ((x, y) for y in range(h))
- # down and to the right
- for x0 in range(-h + 1, w):
- yield ((x0 + y, y) for y in range(max(0, -x0), min(h, w - x0)))
- # down and to the left
- for x0 in range(0, w + h):
- yield ((x0 - y, y) for y in range(max(0, x0 - w + 1), min(h, x0 + 1)))
- def solve(seats: Dict[Tuple[int, int], Set[Tuple[int, int]]], n: int):
- # assume that all seats are free at the beginning.
- free = {seat for seat, visible in seats.items() if len(visible) > n} # after 2 iterations
- while True:
- # visible = {seat, *neighbours}
- changed = set(seat for seat, visible in seats.items() if visible <= free or seat not in free and len(visible - free) > n)
- if not changed:
- break
- free ^= changed
- return len(seats) - len(free)
- if __name__ == '__main__':
- lines = [line.rstrip() for line in sys.stdin]
- seats_part2 = parse_seats(lines)
- 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()}
- print(f'Part 1: {solve(seats_part1, 4)}')
- print(f'Part 2: {solve(seats_part2, 5)}')
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement