day(17): refactor

Signed-off-by: Matej Focko <mfocko@redhat.com>

src/year2021/day17/Day17.kt

@@ -1,17 +1,36 @@
 package year2021.day17
 
 import readInput
+import kotlin.math.absoluteValue
+import kotlin.math.min
 import kotlin.math.max
 import kotlin.math.sign
 
 data class Vector(val x: Int, val y: Int) {
     operator fun plus(other: Vector): Vector = Vector(x + other.x, y + other.y)
+
+    fun updateVelocity(): Vector = Vector(x - x.sign, y - 1)
 }
 
 data class Area(val min: Vector, val max: Vector) {
     fun has(z: Vector): Boolean = (z.x in min.x..max.x) && (z.y in min.y..max.y)
+
+    private fun speeds(getter: (Vector) -> Int): IntRange =
+        max(getter(min).absoluteValue, getter(max).absoluteValue).let { absMax ->
+            (min(absMax * getter(min).sign, 0)..absMax)
+        }
+
+    val xSpeeds: IntRange
+        get() = speeds { it.x }
+
+    val ySpeeds: IntRange
+        get() = speeds { it.y }
+
+    val possibleVectors: Iterable<Vector>
+        get() = xSpeeds.flatMap { vx -> ySpeeds.map { vy -> Vector(vx, vy) } }
 }
 
+// region parsing
 fun getRange(range: String): Pair<Int, Int> = range.split("=")[1].let {
     val (minValue, maxValue) = it.split("..").map(String::toInt)
     minValue to maxValue
@@ -26,51 +45,43 @@ fun String.toArea(): Area {
 
     return Area(Vector(minX, minY), Vector(maxX, maxY))
 }
+// endregion parsing
+
 
 fun findPeak(input: Area, v0: Vector): Int? {
-    fun updateVelocity(v: Vector): Vector = Vector(v.x - v.x.sign, v.y - 1)
-
     var v = v0
     var p = Vector(0, 0)
-
     var maxY = 0
-    var hitTarget = false
 
     while (p.x <= input.max.x && p.y >= input.min.y) {
         p += v
-        v = updateVelocity(v)
+        v = v.updateVelocity()
 
-        hitTarget = hitTarget || input.has(p)
         maxY = max(p.y, maxY)
+        if (input.has(p)) {
+            return maxY
+        }
     }
 
-    return if (hitTarget) maxY else null
+    return null
 }
 
 fun part1(input: Area): Int =
-    (0..input.max.x)
-        .flatMap { x -> (input.min.y..-input.min.y).map { y -> x to y } }
-        .maxOf { (x, y) ->
-            findPeak(input, Vector(x, y)) ?: 0
-        }
-
-// x = v_0x * t + 1/2 * a_x * t^2
-// y = v_0y * t + 1/2 * -1 * t^2
-// v_fy = v_0y + a_y * t => t_up = v_0y / a_y
+    input.possibleVectors.maxOf { (x, y) ->
+        findPeak(input, Vector(x, y)) ?: 0
+    }
 
 fun part2(input: Area): Int =
-    (0..input.max.x)
-        .flatMap { x -> (input.min.y..-input.min.y).map { y -> x to y } }
-        .count { (x, y) ->
-            findPeak(input, Vector(x, y)) != null
-        }
+    input.possibleVectors.count { (x, y) ->
+        findPeak(input, Vector(x, y)) != null
+    }
 
 fun main() {
     val sample = readInput(17, "sample").first().toArea()
     val input = readInput(17, "input").first().toArea()
 
-    println("x: 0..${sample.max.x}, y: ${sample.min.y}..${-sample.min.y}")
-    println("x: 0..${input.max.x}, y: ${input.min.y}..${-input.min.y}")
+    println("x: ${sample.xSpeeds}, y: ${sample.ySpeeds}")
+    println("x: ${input.xSpeeds}, y: ${input.ySpeeds}")
 
     check(part1(sample) == 45)
     println(part1(input))