Further optimise day 4

day04.py

@@ -1,24 +1,27 @@
 import fileinput
 
-#
-# Part 1
-#
-cards = []
-part1 = 0
-for i, line in enumerate(fileinput.input()):
-    card = line.rstrip().split(": ")[1]
-    (winning, owned) = ({int(i) for i in h.split()} for h in card.split(" | "))
+def main():
+    #
+    # Part 1
+    #
+    cards = []
+    part1 = 0
+    for line in fileinput.input():
+        card = line.split(": ", 1)[1]
+        (winning, owned) = ({int(i) for i in h.split()} for h in card.split(" | ", 1))
         count = len(winning & owned)
         cards.append(count)
         part1 += 1 << count - 1 if count else 0
 
-#
-# Part 2
-#
-dp = [0] * len(cards)
-part2 = 0
-for i in range(len(cards)):
+    #
+    # Part 2
+    #
+    dp = [0] * len(cards)
+    part2 = 0
+    for i in range(len(cards)):
         dp[i] = 1 + sum(dp[i - cards[-i - 1] : i])
         part2 += dp[i]
 
-print(part1, part2)
+    print(part1, part2)
+
+main()