Day 11: Part 2 done

Author: tobigrimm

day11/day11.py

@@ -1,81 +1,104 @@
 #!/usr/bin/env python3
 import sys
 from collections import Counter, defaultdict
+from operator import add
 from typing import Tuple, List, Counter as TypeCounter
 
 
 
-def check_neighbors(map, row_num, col_num, symbol) -> int:
+def check_neighbors(seatmap, row_num: int, col_num: int, symbol: str, direct: bool = True) -> int:
     """return the number of symbols around the seat specified
         by row and col"""
     results = 0
-  
-    copy = map.keys()
-    # print("neighbors:", row_num, col_num) 
-    for row in range(row_num-1, row_num+2):
-        for col in range(col_num-1,col_num+2):
-            if (row, col) in copy:
-                #print(row, col)
-                # dont check the seat itself :)
-                if map[(row, col)] == symbol:
-                    if not (row, col) == (row_num, col_num):
-                        results += 1
-    # print("res:", results)
+ 
+    
+    copy = seatmap.keys()
+    if not direct:
+        for row in range(row_num-1, row_num+2):
+            for col in range(col_num-1,col_num+2):
+                if (row, col) in copy:
+                    # dont check the seat itself :)
+                    if seatmap[(row, col)] == symbol:
+                        if not (row, col) == (row_num, col_num):
+                            results += 1
+        # print("res:", results)
+    else:
+        #build direction vectors:
+        for dir in [(0, -1), (0, 1), (-1, 0), (1, 0), (-1, -1), (-1, 1), (1, -1), (1, 1)]:
+            check_key = tuple(map(add, (row_num, col_num), dir))
+            while check_key in copy:
+                if seatmap[check_key] == symbol:
+                    results += 1
+                    break
+                # stop looking if theres an empty seat
+                elif seatmap[check_key] == "L":
+                    break
+
+                check_key = tuple(map(add, check_key, dir))
+
+    # check if im being stupid :)
+    assert(results <= 8)
     return results
 
 
 
 
-def calculate_step(map: List[List[str]]) -> Tuple[bool, List[List[str]]]:
+def calculate_step(seatmap: List[List[str]], lim: int = 4, direct: bool = True) -> Tuple[bool, List[List[str]]]:
     changes = [] 
-    next_map = defaultdict(str)
-    for entry in map.items():
+    next_seatmap = defaultdict(str)
+    for entry in seatmap.items():
         row_num, col_num = entry[0]
         seat = entry[1]
         next_seat = seat
-        num_neighbors = check_neighbors(map, row_num, col_num, "#")
+        num_neighbors = check_neighbors(seatmap, row_num, col_num, "#", direct)
         if seat == "L":
             if num_neighbors == 0:
                 next_seat = "#"
             else:
                 pass
-                #print("Empty seat will stay empty:", row_num, col_num, num_neighbors)
 
-        if seat == "#" and num_neighbors >= 4:
+        if seat == "#" and num_neighbors >= lim:
             next_seat = "L"
         if next_seat != seat:
             changes.append(True)
-            #print("Next: ", row_num, col_num, next_seat, num_neighbors)
-        next_map[(row_num, col_num)] = next_seat
+        next_seatmap[(row_num, col_num)] = next_seat
 
-    return (any(changes), next_map)
+    return (any(changes), next_seatmap)
+
+def calculate_rounds(seatmap, lim: int = 4, direct = False) -> int:
+    round = 0
+    run = True
+    while run:
+        round += 1
+        change, next_seatmap = calculate_step(seatmap, lim, direct)
+        if change:
+            seatmap = next_seatmap
+        else:
+            run = False
+    return list(seatmap.values()).count("#")
+
+def print_nice_map(seatmap):
+    result = ""
+    for row in range(0,100):
+        for col in range(0,100):
+            if (row, col) in seatmap:
+                result += str(seatmap[(row, col)])
+        if (row, 0) in seatmap:
+            result += "\n"
+    return result
 
 if __name__ == "__main__":
-    map = defaultdict(str)
+    seatmap = defaultdict(str)
     with open(sys.argv[1], "r") as infile:
-        # add initial 0 as input jolt
 
-        # TODO change all to defaultdict, use (row, col) as key!
         line = 0
         for nextline in infile:
             for num, char in enumerate(nextline.strip()):
-                map[(line, num)] = char
+                seatmap[(line, num)] = char
             line += 1
 
-    # print(map) 
 
-    round = 0
-    run = True
-    while run:
-        #print(check_neighbors(map, 9, 10, "#"))
-        #print("Round: ", round, "\n Map: ", map)
-        round += 1
-        change, next_map = calculate_step(map)
-        if change:
-            map = next_map
-        else:
-            run = False
-        #if round > 10:
-        #    run = False
-    print("Part1: ", list(map.values()).count("#"))
-    #print("Part2: ", get_variants(jolts))
+    part1 = calculate_rounds(seatmap.copy())
+    print("Part1: ", part1)
+    part2 = calculate_rounds(seatmap.copy(), 5, True)
+    print("Part2: ", part2)