Day 11
[advent-of-code-2020.git] / day11 / seats.py
1 #!/usr/bin/python3
2
3 import sys
4 import copy
5
6 base_states=[]
7 states=[]
8
9 filename=(len(sys.argv) > 1) and sys.argv[1] or "37_26_seats.txt"
10
11 base_states=[list(s.rstrip()) for s in open(filename, "r")]
12 states=copy.deepcopy(base_states)
13
14 def get_adjacent_seats(states, y, x, seat_map):
15     occupied=0
16     start_y=y-1
17     if y == 0:
18         start_y=0
19
20     end_y=y+1
21     if y >= len(states)-1:
22         end_y=y
23
24     start_x=x-1
25     if x == 0:
26         start_x=0
27
28     end_x=x+1
29     if x >= len(states[y])-1:
30         end_x=x
31
32     for check_y in range(start_y, end_y+1):
33         for check_x in range(start_x, end_x+1):
34             if check_y == y and check_x == x:
35                 continue
36             if states[check_y][check_x] == "#":
37                 occupied+=1
38
39     return occupied
40
41 def get_nearest_seats(states, y, x, seat_map):
42     occupied=0
43     for seat in seat_map[y][x]:
44         if states[seat[0]][seat[1]] == "#":
45             occupied+=1
46
47     return occupied
48
49 def calc_nearest_seats(states):
50     seats={}
51     for y,row in enumerate(states):
52         for x,seat in enumerate(row):
53             # only do anything if the seat is a seat
54             if seat == '.':
55                 continue
56             daseats={}
57             for check_y,other_row in enumerate(states):
58                 for check_x,other_seat in enumerate(other_row):
59                     # only do anything if the seat is a seat
60                     if other_seat == '.':
61                         continue
62                     # skip outselves
63                     if check_x == x and check_y == y:
64                         continue
65                     if abs(check_x - x) == abs(check_y - y):
66                         # this is one of the four diagonals
67                         if check_x > x and check_y > y:
68                             # down and right, getting further away
69                             # so if we've got a value ignore, other
70                             # wise add it
71                             if 'dr' not in daseats:
72                                 daseats['dr']=[check_y,check_x]
73                         elif check_x > x and check_y < y:
74                             # up and right, getting closer
75                             daseats['ur']=[check_y,check_x]
76                         elif check_x < x and check_y < y:
77                             # up and left, getting closer, so always replace
78                             daseats['ul']=[check_y,check_x]
79                         elif check_x < x and check_y > y:
80                             # down and left getting further away
81                             if 'dl' not in daseats:
82                                 daseats['dl']=[check_y,check_x]
83                     elif check_x > x and check_y == y:
84                         # heading right, getting further away
85                         # so if we've got a value ingore, other
86                         # wise add it
87                         if 'r' not in daseats:
88                             daseats['r']=[check_y,check_x]
89                     elif check_x < x and check_y == y:
90                         # heading left, getting closer
91                         daseats['l']=[check_y,check_x]
92                     elif check_y > y and check_x == x:
93                         # heading down, getting further away
94                         if 'd' not in daseats:
95                             daseats['d']=[check_y,check_x]
96                     elif check_y < y and check_x == x:
97                         # heading up
98                         daseats['u']=[check_y,check_x]
99             if not y in seats:
100                 seats[y]={x: [seat for seat in daseats.values()]}
101             seats[y][x]=[seats for seats in daseats.values()]
102
103     return seats
104
105
106 def apply_rules(states, check=get_adjacent_seats, max_people=4, seat_map=None):
107     new_states=copy.deepcopy(states)
108
109     for y in range(0, len(states)):
110         for x in range(0, len(states[0])):
111             if states[y][x] == ".":
112                 continue
113             adj_seats=check(states, y, x, seat_map)
114             if states[y][x] == 'L' and adj_seats == 0:
115                 new_states[y][x]='#'
116             elif states[y][x] == '#' and adj_seats >= max_people:
117                 new_states[y][x]='L'
118
119     return new_states
120
121 def get_seats(states):
122     return ''.join([''.join(s) for s in states])
123
124 def print_layout(states, indent_string=""):
125     for line in states:
126         print(indent_string, ''.join(line))
127
128     print()
129
130 previous_seats=get_seats(states)
131 states=apply_rules(states)
132 seats=get_seats(states)
133
134 while seats!=previous_seats:
135     states=apply_rules(states)
136     previous_seats=seats
137     seats=get_seats(states)
138
139 print("Part 1:")
140 print("  Final state:")
141 print_layout(states, "  ")
142 print("  Occupied: {}".format(seats.count('#')))
143
144 # now do part 2!
145 states=copy.deepcopy(base_states)
146 seat_map=calc_nearest_seats(states)
147
148 previous_seats=get_seats(states)
149 states=apply_rules(states,get_nearest_seats,5,seat_map)
150 seats=get_seats(states)
151
152 while seats!=previous_seats:
153     states=apply_rules(states,get_nearest_seats,5,seat_map)
154     previous_seats=seats
155     seats=get_seats(states)
156
157 print("Part 2:")
158 print("  Final state:")
159 print_layout(states, "  ")
160 print("  Occupied: {}".format(seats.count('#')))