diff --git a/app/3.hs b/app/3.hs
index 8fc1c50..2cd9f80 100644
--- a/app/3.hs
+++ b/app/3.hs
@@ -7,51 +7,32 @@ import Data.Char (digitToInt)
 import Control.Applicative
 import Data.Maybe
 import Data.Function (on)
-import Debug.Trace (trace)
 import Data.Ord
 
-
 parse :: String -> [[Int]]
 parse =  map (digitToInt <$>) . lines
 
-solve1 :: [[Int]] -> Int
-solve1 = sum . map go
+solve :: Int -> [[Int]] -> Int
+solve n = sum . mapMaybe (\xs -> go n [] xs xs)
   where
-    go :: [Int] -> Int
-    go xs = let m = maximum xs
-                (Just mi) = elemIndex m xs
-                xsafter = drop (mi+1) xs
-                mafter = maximum xsafter
-                m2 = maximum (delete m xs)
-            in if xsafter /= [] then
-                 m  * 10 + mafter else
-                 m2 * 10 + m
-
-solve2 :: [[Int]] -> Int
-solve2 = sum . mapMaybe (go)
-  where
-    go :: [Int] -> Maybe Int
-    go ys = go' 12 [] ys
-      where
-        go' :: Int -> [(Int,Int)] -> [Int] -> Maybe Int
-        go' 0     acc _  =
-          (pure . foldl ((\a b -> a * 10 + b)) 0 . map snd $ sortBy (compare `on` fst) acc)
-        go' (n+1) acc xs
-          | length xs < n+1 = Nothing
-          | otherwise =
-                let
-                    snMax = sortBy (comparing Down) . fst $ foldl (\(a,s) _ ->
-                                     let curr = maximum s in
-                                     (curr:a,delete curr s)) ([],xs) [1..n+1]
-                in
-                  asum (map (\m ->
-                               let (Just mi) = elemIndex m xs
-                                   xsafter = drop (mi+1) xs
-                                   (Just mio) = elemIndex m ys in
-                                  go' n ((mio,m):acc) xsafter)
-                         snMax)
+    go :: Int -> [(Int,Int)] -> [Int] -> [Int] -> Maybe Int
+    go 0     acc _ _  =
+      pure . foldl (\a b -> a * 10 + b) 0 . map snd $ sortBy (compare `on` fst) acc
+    go (n+1) acc ys xs
+      | length xs < n+1 = Nothing
+      | otherwise =
+            let snMax = sortBy (comparing Down) . fst $ foldl (\(a,s) _ ->
+                                 let curr = maximum s in
+                                 (curr:a,delete curr s)) ([],xs) [1..n+1]
+            in
+              asum (map (\m ->
+                           let (Just mi) = elemIndex m xs
+                               xsafter = drop (mi+1) xs
+                               (Just mio) = elemIndex m ys in
+                              go n ((mio,m):acc) (replicate (mio + 1) 0 ++ drop (mio + 1) ys) xsafter)
+                     snMax)
 
 main :: IO ()
 main = readFile "inputs/3" <&> parse >>= \i ->
-  print (solve1 i) >>
-  print (solve2 i)
+  print (solve  2 i) >>
+  print (solve 12 i)