diff --git a/src/main.rs b/src/main.rs
index 48165d4..ace2b46 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,4 +1,6 @@
 #![feature(test)]
+
+use std::collections::BinaryHeap;
 extern crate test;
 
 const INPUTS: [&'static str; 2] = [
@@ -16,36 +18,69 @@ fn parse_input(input: &'static str) -> Vec<Vec<u8>> {
 
 const NEIGHBOURS: [[i32; 2]; 4] = [[-1, 0], [1, 0], [0, -1], [0, 1]];
 
-fn solution(input: Vec<Vec<u8>>) -> i32 {
-    let mut answer = 0;
+fn solution(mut input: Vec<Vec<u8>>) -> u64 {
+    let m = input.len() as i32;
+    let n = input[0].len() as i32;
 
-    let m = input.len();
-    let n = input[0].len();
+    let mut heap = BinaryHeap::with_capacity(3);
 
     for i in 0..m {
         for j in 0..n {
             let mut is_minimum = true;
 
             for &neighbour in NEIGHBOURS.iter() {
-                let x = i as i32 + neighbour[0];
-                let y = j as i32 + neighbour[1];
+                let x = i + neighbour[0];
+                let y = j + neighbour[1];
 
-                if x < 0 || y < 0 || x >= m as i32 || y >= n as i32 {
+                if x < 0 || y < 0 || x >= m || y >= n {
                     continue;
                 }
 
-                if input[i][j] >= input[x as usize][y as usize] {
+                if input[i as usize][j as usize] >= input[x as usize][y as usize] {
                     is_minimum = false;
+                    break;
                 }
             }
 
             if is_minimum {
-                answer += input[i][j] as i32 + 1;
+                let basin_size = dfs(&mut input, i, j);
+
+                heap.push(basin_size);
             }
         }
     }
 
-    answer
+    heap.into_iter().take(3).fold(1, |a, x| a * x)
+}
+
+const VISITED: u8 = 1 << 7;
+
+fn dfs(grid: &mut [Vec<u8>], px: i32, py: i32) -> u64 {
+    if grid[px as usize][py as usize] == VISITED {
+        return 0;
+    } else {
+        grid[px as usize][py as usize] = VISITED;
+    }
+
+    let mut ldepth = 0;
+    for neigh in NEIGHBOURS.iter() {
+        let x = px as i32 + neigh[0];
+        let y = py as i32 + neigh[1];
+
+        if x < 0
+            || y < 0
+            || x >= grid.len() as i32
+            || y >= grid[0].len() as i32
+            || grid[x as usize][y as usize] == 9
+            || grid[x as usize][y as usize] == VISITED
+        {
+            continue;
+        }
+
+        ldepth += dfs(grid, x, y);
+    }
+
+    ldepth + 1
 }
 
 fn main() {