{-# Language QuasiQuotes, TemplateHaskell #-}
module Main where
import Advent
import Data.List
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.Set as Set
data E = Egain | Elose deriving Int -> E -> ShowS
[E] -> ShowS
E -> String
(Int -> E -> ShowS) -> (E -> String) -> ([E] -> ShowS) -> Show E
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> E -> ShowS
showsPrec :: Int -> E -> ShowS
$cshow :: E -> String
show :: E -> String
$cshowList :: [E] -> ShowS
showList :: [E] -> ShowS
Show
stageTH
data Edge = Edge String String
deriving (Edge -> Edge -> Bool
(Edge -> Edge -> Bool) -> (Edge -> Edge -> Bool) -> Eq Edge
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Edge -> Edge -> Bool
== :: Edge -> Edge -> Bool
$c/= :: Edge -> Edge -> Bool
/= :: Edge -> Edge -> Bool
Eq, Eq Edge
Eq Edge =>
(Edge -> Edge -> Ordering)
-> (Edge -> Edge -> Bool)
-> (Edge -> Edge -> Bool)
-> (Edge -> Edge -> Bool)
-> (Edge -> Edge -> Bool)
-> (Edge -> Edge -> Edge)
-> (Edge -> Edge -> Edge)
-> Ord Edge
Edge -> Edge -> Bool
Edge -> Edge -> Ordering
Edge -> Edge -> Edge
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Edge -> Edge -> Ordering
compare :: Edge -> Edge -> Ordering
$c< :: Edge -> Edge -> Bool
< :: Edge -> Edge -> Bool
$c<= :: Edge -> Edge -> Bool
<= :: Edge -> Edge -> Bool
$c> :: Edge -> Edge -> Bool
> :: Edge -> Edge -> Bool
$c>= :: Edge -> Edge -> Bool
>= :: Edge -> Edge -> Bool
$cmax :: Edge -> Edge -> Edge
max :: Edge -> Edge -> Edge
$cmin :: Edge -> Edge -> Edge
min :: Edge -> Edge -> Edge
Ord)
edge :: String -> String -> Edge
edge :: String -> String -> Edge
edge String
a String
b
| String
a String -> String -> Bool
forall a. Ord a => a -> a -> Bool
< String
b = String -> String -> Edge
Edge String
a String
b
| Bool
otherwise = String -> String -> Edge
Edge String
b String
a
main :: IO ()
IO ()
main =
do [(String, E, Int, String)]
input <- [format|2015 13 (%s would @E %u happiness units by sitting next to %s.%n)*|]
let graph :: Map Edge Int
graph = (Int -> Int -> Int) -> [(Edge, Int)] -> Map Edge Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(String -> String -> Edge
edge String
x String
y, case E
e of E
Egain -> Int
v; E
Elose -> -Int
v) | (String
x,E
e, Int
v,String
y) <- [(String, E, Int, String)]
input]
let people1 :: [String]
people1 = [String] -> [String]
forall a. Ord a => [a] -> [a]
uniques [String
z | (String
x,E
_,Int
_,String
y) <- [(String, E, Int, String)]
input, String
z <- [String
x,String
y]]
Int -> IO ()
forall a. Show a => a -> IO ()
print (Map Edge Int -> [String] -> Int
maximumHappiness Map Edge Int
graph [String]
people1)
let people2 :: [String]
people2 = String
"" String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
people1
Int -> IO ()
forall a. Show a => a -> IO ()
print (Map Edge Int -> [String] -> Int
maximumHappiness Map Edge Int
graph [String]
people2)
neighbors :: [String] -> [Edge]
neighbors :: [String] -> [Edge]
neighbors [] = []
neighbors (String
x:[String]
xs) = (String -> String -> Edge) -> [String] -> [String] -> [Edge]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith String -> String -> Edge
edge (String
xString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
xs) ([String]
xs [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [String
x])
maximumHappiness ::
Map Edge Int ->
[String] ->
Int
maximumHappiness :: Map Edge Int -> [String] -> Int
maximumHappiness Map Edge Int
relationships [String]
people = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([String] -> Int
score ([String] -> Int) -> [[String]] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String] -> [[String]]
forall a. [a] -> [[a]]
permutations [String]
people)
where
score :: [String] -> Int
score [String]
xs = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Edge -> Map Edge Int -> Int
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Int
0 Edge
e Map Edge Int
relationships | Edge
e <- [String] -> [Edge]
neighbors [String]
xs]
uniques :: Ord a => [a] -> [a]
uniques :: forall a. Ord a => [a] -> [a]
uniques = Set a -> [a]
forall a. Set a -> [a]
Set.toList (Set a -> [a]) -> ([a] -> Set a) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList