diff --git a/src/year2021/day09/Day09.kt b/src/year2021/day09/Day09.kt
index a484233..924c02d 100644
--- a/src/year2021/day09/Day09.kt
+++ b/src/year2021/day09/Day09.kt
@@ -3,39 +3,34 @@ package year2021.day09
 import readInput
 import kotlin.math.max
 
-fun getRiskLevel(input: List<List<Int>>, y: Int, x: Int): Int = listOf(
+fun validAdjacentIndices(input: List<List<Int>>, y: Int, x: Int): Iterable<Pair<Int, Int>> = listOf(
     Pair(y - 1, x),
     Pair(y + 1, x),
     Pair(y, x + 1),
     Pair(y, x - 1)
 ).filter { (y0, x0) -> y0 >= 0 && y0 < input.size && x0 >= 0 && x0 < input[y0].size }
+
+fun getRiskLevel(input: List<List<Int>>, y: Int, x: Int): Int = validAdjacentIndices(input, y, x)
     .map { (y0, x0) -> input[y0][x0] }
-    .all { it > input[y][x] }.let { isLowPoint -> if (isLowPoint) input[y][x] + 1 else 0 }
+    .all { it > input[y][x] }
+    .let { isLowPoint -> if (isLowPoint) input[y][x] + 1 else 0 }
 
 fun part1(input: List<List<Int>>): Int = input
     .mapIndexed { y, row ->
-        row.mapIndexed { x, col ->
+        row.indices.map { x ->
             getRiskLevel(input, y, x)
         }
     }.sumOf { row -> row.sum() }
 
-fun dfs(input: MutableList<MutableList<Int>>, y: Int, x: Int): Int {
-    if (input[y][x] != 0) {
-        return 0
+fun dfs(input: MutableList<MutableList<Int>>, y: Int, x: Int): Int = when (input[y][x]) {
+    0 -> {
+        // mark as visited
+        input[y][x] = -1
+
+        // check adjacent
+        1 + validAdjacentIndices(input, y, x).sumOf { (y0, x0) -> dfs(input, y0, x0) }
     }
-
-    input[y][x] = -1
-
-    var total = 1
-    listOf(
-        Pair(y - 1, x),
-        Pair(y + 1, x),
-        Pair(y, x + 1),
-        Pair(y, x - 1)
-    ).filter { (y0, x0) -> y0 >= 0 && y0 < input.size && x0 >= 0 && x0 < input[y0].size }
-        .forEach { (y0, x0) -> total += dfs(input, y0, x0) }
-
-    return total;
+    else -> 0
 }
 
 fun part2(input: List<List<Int>>): Int {
@@ -56,8 +51,6 @@ fun part2(input: List<List<Int>>): Int {
     }
 
     sizes.sortDescending()
-    println(sizes)
-
     return sizes.take(3).fold(1) { product, n -> n * product }
 }