fix bug with ties and refactor

2021/py/day03.py

@@ -1,4 +1,5 @@
 from aocd import get_data
+import operator
 import pandas as pd
 
 
@@ -9,32 +10,28 @@ def parse_data():
 
 
 def part_one(data):
-    gamma = ''.join([data[i].value_counts().idxmax() for i in data.columns])
-    epsilon = ''.join([data[i].value_counts().idxmin() for i in data.columns])
+    gamma = "".join([data[i].value_counts().idxmax() for i in data.columns])
+    epsilon = "".join([data[i].value_counts().idxmin() for i in data.columns])
     return int(gamma, 2) * int(epsilon, 2)
 
 
-def oxygen_generator(data):
+def get_rating(data, bit_criteria="O2"):
+    BIT_CRITERIA = {"O2": operator.eq, "CO2": operator.ne}
     df = data.copy()
+    op = BIT_CRITERIA[bit_criteria]
     for i in df.columns:
-        df = df[df[i] == df[i].value_counts().idxmax()]
+        col_max = df[i].value_counts().idxmax()
+        if col_max == df[i].value_counts().idxmin():
+            col_max = "1"
+        df = df[op(df[i], col_max)]
         if len(df) == 1:
             break
-    return ''.join(df.values[0])
-
-
-def co2_scrubber(data):
-    df = data.copy()
-    for i in data.columns:
-        df = df[df[i] != df[i].value_counts().idxmax()]
-        if len(df) == 1:
-            break
-    return ''.join(df.values[0])
+    return "".join(df.values[0])
 
 
 def part_two(data):
-    oxygen_rating = oxygen_generator(data)
-    co2_rating = co2_scrubber(data)
+    oxygen_rating = get_rating(data, "O2")
+    co2_rating = get_rating(data, "CO2")
     return int(oxygen_rating, 2) * int(co2_rating, 2)
 
 
@@ -44,4 +41,3 @@ def part_two(data):
     print(part_one(data))
     print("*" * 10, "Part Two", "*" * 10)
     print(part_two(data))
-