day 6 fancy solution + cleanup

day06/sol.py

@@ -1,11 +1,12 @@
+from math import sqrt
 sample = {
-    "time": [7, 15,  30],
+    "time":     [7, 15,  30],
     "distance": [9, 40, 200],
 }
 
 input = {
-"time":      [  48,     87,     69,     81,],
+    "time":     [   48,     87,     69,     81,],
-"distance": [  255,   1288,   1117,   1623,],
+    "distance": [  255,   1288,   1117,   1623,],
 }
 
 
@@ -21,18 +22,54 @@ def solve(data):
             d = (t-i)*v
             if d > best:
                 ways += 1
+
         part1 *= ways
 
     return part1
 
+
+def fancy_solve(data):
+    time = data["time"]
+    dist = data["distance"]
+
+    for t, best in zip(time, dist):
+        # the distance traveled if the boat is released at time i
+        # is equal to (t-i)*i
+        # we want to know the range of values of i for which this
+        # exceeds the best time
+        #   (t-i)*i > best
+        # a little rearranging gets us the quadratic equation
+        #   i^2 - ti + best <= 0
+        # which we can determine the crossing points for
+        # using the quadratic formula (the good one, not the
+        # one you learned in school)
+        h = t/2
+        s = sqrt(h*h - best)
+        #print(h-s, h+s)
+
+        # in general these will not be at integer values,
+        # so we probe the floor and ceiling of each crossing
+        # point to determine exactly where the condition is met
+        a = int(h-s)
+        b = int(h+s)
+        if (t-a)*a <= best:
+            a += 1
+        if (t-b)*b > best:
+            b += 1
+
+        ways = b - a
+        return ways
+
+
 print(solve(sample))
 print(solve(input))
 
 def part2(data):
-    return {"time": [int("".join(str(x) for x in data["time"]))],
+    return {"time":     [int("".join(str(x) for x in data["time"]))],
             "distance": [int("".join(str(x) for x in data["distance"]))]}
 
 print(solve(part2(sample)))
-print(solve(part2(input)))
+print(fancy_solve(part2(sample)))
+print(fancy_solve(part2(input)))