Faster non-recursive implementation with maths

ephemient · Dec 9, 2023 · a2b9a4c · a2b9a4c
1 parent 6a540fc
commit a2b9a4c
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 25 deletions.
diff --git a/hs/src/Day9.hs b/hs/src/Day9.hs
@@ -5,13 +5,16 @@ Description:    <https://adventofcode.com/2023/day/9 Day 9: Mirage Maintenance>
 module Day9 (part1, part2) where
 
 import Common (readEntire, readMany)
+import Data.List (foldl', scanl')
 import Data.Text (Text)
 import qualified Data.Text as T (lines)
 import qualified Data.Text.Read as T (decimal, signed)
 
-predict :: (Num a, Eq a) => [a] -> a
-predict xs | all (== 0) xs = 0
-predict xs@(_:_) = last xs + predict (zipWith subtract xs $ tail xs)
+binom :: Int -> [Int]
+binom n = scanl' f 1 [1..n] where f k i = k * (n + 1 - i) `div` i
+
+predict :: [Int] -> Int
+predict xs = foldl' subtract 0 $ zipWith (*) xs $ binom $ length xs
 
 part1, part2 :: Text -> Either String Int
 part1 = fmap sum .  mapM (fmap predict . readEntire (readMany $ T.signed T.decimal)) .  T.lines

diff --git a/kt/aoc2023-lib/src/commonMain/kotlin/com/github/ephemient/aoc2023/Day9.kt b/kt/aoc2023-lib/src/commonMain/kotlin/com/github/ephemient/aoc2023/Day9.kt
@@ -10,10 +10,14 @@ class Day9(input: String) {
     fun part2(): Int = nums.sumOf { it.asReversed().extrapolate() }
 
     companion object {
-        private fun List<Int>.extrapolate(): Int = if (any { it != 0 }) {
-            last() + zipWithNext { a, b -> b - a }.extrapolate()
-        } else {
-            0
+        private fun List<Int>.extrapolate(): Int {
+            var c = 1
+            var s = 0
+            for ((i, x) in withIndex()) {
+                s = c * x - s
+                c = c * (size - i) / (i + 1)
+            }
+            return s
         }
     }
 }
diff --git a/py/aoc2023/day9.py b/py/aoc2023/day9.py
@@ -10,9 +10,11 @@
 
 
 def _extrapolate(nums):
-    if any(nums):
-        return nums[-1] + _extrapolate([y - x for x, y in zip(nums, nums[1:])])
-    return 0
+    c, s, n = 1, 0, len(nums)
+    for i, x in enumerate(nums):
+        s = c * x - s
+        c = c * (n - i) // (i + 1)
+    return s
 
 
 def part1(data):

diff --git a/rs/src/day9.rs b/rs/src/day9.rs
@@ -1,19 +1,13 @@
-fn predict(nums: &[i32]) -> i32 {
-    if nums.iter().any(|&x| x != 0) {
-        nums.last().unwrap()
-            + predict(
-                &nums
-                    .iter()
-                    .zip(nums.iter().skip(1))
-                    .map(|(x, y)| y - x)
-                    .collect::<Vec<_>>(),
-            )
-    } else {
-        0
-    }
+fn predict(nums: &[i32]) -> i64 {
+    nums.iter()
+        .enumerate()
+        .fold((1, 0), |(c, s), (i, x)| {
+            (c * (nums.len() - i) / (i + 1), c as i64 * *x as i64 - s)
+        })
+        .1
 }
 
-pub fn part1(data: &str) -> i32 {
+pub fn part1(data: &str) -> i64 {
     data.lines()
         .map(|line| {
             predict(
@@ -26,7 +20,7 @@ pub fn part1(data: &str) -> i32 {
         .sum()
 }
 
-pub fn part2(data: &str) -> i32 {
+pub fn part2(data: &str) -> i64 {
     data.lines()
         .map(|line| {
             predict(