[py] Port 2016_04, kinda

py/aoc/y2016/d04.py

@@ -0,0 +1,83 @@
+import re
+
+# PART 1
+
+ROOM_RE = r"(.*)-(\d+)\[(.{5})\]"
+
+
+def load_rooms(inputstr):
+    rooms = []
+    for line in inputstr.splitlines():
+        match = re.fullmatch(ROOM_RE, line)
+        name, sector, checksum = match.groups()
+        name = name.split("-")
+        sector = int(sector)
+        rooms.append((name, sector, checksum))
+    return rooms
+
+
+def checksum(name):
+    name = "".join(name)
+    freqs = {}
+    for char in name:
+        freqs[char] = freqs.get(char, 0) + 1
+
+    top = freqs.items()
+    top = sorted(top, key=lambda x: x[0])
+    top = sorted(top, key=lambda x: x[1], reverse=True)
+    top = list(top)[:5]
+    top = map(lambda x: x[0], top)
+    return "".join(top)
+
+
+def sum_correct_sectors(rooms):
+    total = 0
+    for name, sector, chk in rooms:
+        if checksum(name) == chk:
+            total += sector
+    return total
+
+
+# PART 2
+
+
+def shift(s):
+    new_s = []
+    for c in s:
+        if ord("a") <= ord(c) <= ord("z"):
+            new_c = chr(
+                ord("a") + ((ord(c) - ord("a") + 1) % (ord("z") - ord("a") + 1))
+            )
+            new_s.append(new_c)
+        elif ord("A") <= ord(c) <= ord("Z"):
+            new_c = chr(
+                ord("A") + ((ord(c) - ord("A") + 1) % (ord("Z") - ord("A") + 1))
+            )
+            new_s.append(new_c)
+        else:
+            new_s.append("?")
+    return "".join(new_s)
+
+
+def decrypt_name(name, sector):
+    for _ in range(sector % 26):
+        name = map(shift, name)
+    return " ".join(name)
+
+
+def decrypt_names(rooms, filename):
+    with open(filename, "w") as f:
+        for name, sector, chk in rooms:
+            if checksum(name) == chk:
+                name = decrypt_name(name, sector)
+                line = f"{name}-{sector}[{chk}]\n"
+                f.write(line)
+
+
+def solve(inputstr):
+    rooms = load_rooms(inputstr)
+    total = sum_correct_sectors(rooms)
+    print(f"Part 1: {total}")
+    # filename_2 = inputstr + ".decrypted"
+    # decrypt_names(rooms, filename_2)
+    # print(f"Part 2: see {filename_2} (search for 'northpole object storage')")