day 6 fancy solution + cleanup

day06/sol.py

@@ -1,3 +1,4 @@
+from math import sqrt
 sample = {
     "time":     [7, 15,  30],
     "distance": [9, 40, 200],
@@ -21,10 +22,45 @@ 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))
 
@@ -33,6 +69,7 @@ def part2(data):
             "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)))