{-# Language QuasiQuotes, ImportQualifiedPost, DeriveTraversable #-}
module Main
(
main
, Node(Node)
, Summary(Summary)
, topName
, summarize
, computeCorrection
, corrections
, OneChangeT(OCT)
, change
) where
import Advent (format, pickOne, same)
import Advent.Fix (Fix, cataM, anaFromMap)
import Control.Applicative (Alternative(empty,(<|>),some))
import Control.Monad (MonadPlus, ap, liftM)
import Data.Foldable (asum)
import Data.Functor.Classes (Show1(liftShowsPrec))
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (fromMaybe, listToMaybe)
import Data.Set (Set)
import Data.Set qualified as Set
data Node a = Node !Int [a]
deriving (Int -> Node a -> ShowS
[Node a] -> ShowS
Node a -> [Char]
(Int -> Node a -> ShowS)
-> (Node a -> [Char]) -> ([Node a] -> ShowS) -> Show (Node a)
forall a. Show a => Int -> Node a -> ShowS
forall a. Show a => [Node a] -> ShowS
forall a. Show a => Node a -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Node a -> ShowS
showsPrec :: Int -> Node a -> ShowS
$cshow :: forall a. Show a => Node a -> [Char]
show :: Node a -> [Char]
$cshowList :: forall a. Show a => [Node a] -> ShowS
showList :: [Node a] -> ShowS
Show, (forall a b. (a -> b) -> Node a -> Node b)
-> (forall a b. a -> Node b -> Node a) -> Functor Node
forall a b. a -> Node b -> Node a
forall a b. (a -> b) -> Node a -> Node b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Node a -> Node b
fmap :: forall a b. (a -> b) -> Node a -> Node b
$c<$ :: forall a b. a -> Node b -> Node a
<$ :: forall a b. a -> Node b -> Node a
Functor, (forall m. Monoid m => Node m -> m)
-> (forall m a. Monoid m => (a -> m) -> Node a -> m)
-> (forall m a. Monoid m => (a -> m) -> Node a -> m)
-> (forall a b. (a -> b -> b) -> b -> Node a -> b)
-> (forall a b. (a -> b -> b) -> b -> Node a -> b)
-> (forall b a. (b -> a -> b) -> b -> Node a -> b)
-> (forall b a. (b -> a -> b) -> b -> Node a -> b)
-> (forall a. (a -> a -> a) -> Node a -> a)
-> (forall a. (a -> a -> a) -> Node a -> a)
-> (forall a. Node a -> [a])
-> (forall a. Node a -> Bool)
-> (forall a. Node a -> Int)
-> (forall a. Eq a => a -> Node a -> Bool)
-> (forall a. Ord a => Node a -> a)
-> (forall a. Ord a => Node a -> a)
-> (forall a. Num a => Node a -> a)
-> (forall a. Num a => Node a -> a)
-> Foldable Node
forall a. Eq a => a -> Node a -> Bool
forall a. Num a => Node a -> a
forall a. Ord a => Node a -> a
forall m. Monoid m => Node m -> m
forall a. Node a -> Bool
forall a. Node a -> Int
forall a. Node a -> [a]
forall a. (a -> a -> a) -> Node a -> a
forall m a. Monoid m => (a -> m) -> Node a -> m
forall b a. (b -> a -> b) -> b -> Node a -> b
forall a b. (a -> b -> b) -> b -> Node a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => Node m -> m
fold :: forall m. Monoid m => Node m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Node a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Node a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Node a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> Node a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> Node a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Node a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Node a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Node a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Node a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Node a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Node a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Node a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> Node a -> a
foldr1 :: forall a. (a -> a -> a) -> Node a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Node a -> a
foldl1 :: forall a. (a -> a -> a) -> Node a -> a
$ctoList :: forall a. Node a -> [a]
toList :: forall a. Node a -> [a]
$cnull :: forall a. Node a -> Bool
null :: forall a. Node a -> Bool
$clength :: forall a. Node a -> Int
length :: forall a. Node a -> Int
$celem :: forall a. Eq a => a -> Node a -> Bool
elem :: forall a. Eq a => a -> Node a -> Bool
$cmaximum :: forall a. Ord a => Node a -> a
maximum :: forall a. Ord a => Node a -> a
$cminimum :: forall a. Ord a => Node a -> a
minimum :: forall a. Ord a => Node a -> a
$csum :: forall a. Num a => Node a -> a
sum :: forall a. Num a => Node a -> a
$cproduct :: forall a. Num a => Node a -> a
product :: forall a. Num a => Node a -> a
Foldable, Functor Node
Foldable Node
(Functor Node, Foldable Node) =>
(forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Node a -> f (Node b))
-> (forall (f :: * -> *) a.
Applicative f =>
Node (f a) -> f (Node a))
-> (forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Node a -> m (Node b))
-> (forall (m :: * -> *) a. Monad m => Node (m a) -> m (Node a))
-> Traversable Node
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Node (m a) -> m (Node a)
forall (f :: * -> *) a. Applicative f => Node (f a) -> f (Node a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Node a -> m (Node b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Node a -> f (Node b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Node a -> f (Node b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Node a -> f (Node b)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Node (f a) -> f (Node a)
sequenceA :: forall (f :: * -> *) a. Applicative f => Node (f a) -> f (Node a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Node a -> m (Node b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Node a -> m (Node b)
$csequence :: forall (m :: * -> *) a. Monad m => Node (m a) -> m (Node a)
sequence :: forall (m :: * -> *) a. Monad m => Node (m a) -> m (Node a)
Traversable)
instance Show1 Node where
liftShowsPrec :: forall a.
(Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> Node a -> ShowS
liftShowsPrec Int -> a -> ShowS
_ [a] -> ShowS
s Int
p (Node Int
x [a]
y) =
Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
11)
([Char] -> ShowS
showString [Char]
"Node " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 Int
x ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> ShowS
showChar Char
' ' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> ShowS
s [a]
y)
main :: IO ()
IO ()
main =
do [([Char], Int, Maybe [[Char]])]
input <- [format|2017 7 (%s %(%u%)(| -> %s&(, ))%n)*|]
let nodes :: Map [Char] (Node [Char])
nodes = [([Char], Node [Char])] -> Map [Char] (Node [Char])
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [([Char]
n, Int -> [[Char]] -> Node [Char]
forall a. Int -> [a] -> Node a
Node Int
c ([[Char]] -> Maybe [[Char]] -> [[Char]]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [[Char]]
xs)) | ([Char]
n, Int
c, Maybe [[Char]]
xs) <- [([Char], Int, Maybe [[Char]])]
input]
let top :: [Char]
top = Map [Char] (Node [Char]) -> [Char]
forall name. Ord name => Map name (Node name) -> name
topName Map [Char] (Node [Char])
nodes
[Char] -> IO ()
putStrLn [Char]
top
Maybe Int -> IO ()
forall a. Show a => a -> IO ()
print (Fix Node -> Maybe Int
computeCorrection (Map [Char] (Node [Char]) -> [Char] -> Fix Node
forall k (f :: * -> *).
(Ord k, Functor f) =>
Map k (f k) -> k -> Fix f
anaFromMap Map [Char] (Node [Char])
nodes [Char]
top))
topName :: Ord name => Map name (Node name) -> name
topName :: forall name. Ord name => Map name (Node name) -> name
topName Map name (Node name)
m = Set name -> name
forall a. Set a -> a
Set.findMin
(Set name -> name) -> Set name -> name
forall a b. (a -> b) -> a -> b
$ Set name -> Set name -> Set name
forall a. Ord a => Set a -> Set a -> Set a
Set.difference
(Map name (Node name) -> Set name
forall k a. Map k a -> Set k
Map.keysSet Map name (Node name)
m)
([name] -> Set name
forall a. Ord a => [a] -> Set a
Set.fromList ((Node name -> [name]) -> Map name (Node name) -> [name]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(Node Int
_ [name]
xs) -> [name]
xs) Map name (Node name)
m))
data Summary = Summary !Int !Int
deriving Int -> Summary -> ShowS
[Summary] -> ShowS
Summary -> [Char]
(Int -> Summary -> ShowS)
-> (Summary -> [Char]) -> ([Summary] -> ShowS) -> Show Summary
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Summary -> ShowS
showsPrec :: Int -> Summary -> ShowS
$cshow :: Summary -> [Char]
show :: Summary -> [Char]
$cshowList :: [Summary] -> ShowS
showList :: [Summary] -> ShowS
Show
summarize :: Fix Node -> OneChangeT Int [] Summary
summarize :: Fix Node -> OneChangeT Int [] Summary
summarize = (Node Summary -> OneChangeT Int [] Summary)
-> Fix Node -> OneChangeT Int [] Summary
forall (m :: * -> *) (t :: * -> *) a.
(Monad m, Traversable t) =>
(t a -> m a) -> Fix t -> m a
cataM ((Node Summary -> OneChangeT Int [] Summary)
-> Fix Node -> OneChangeT Int [] Summary)
-> (Node Summary -> OneChangeT Int [] Summary)
-> Fix Node
-> OneChangeT Int [] Summary
forall a b. (a -> b) -> a -> b
$ \(Node Int
n [Summary]
xs) ->
if [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => t a -> Bool
same [ Int
w | Summary Int
_ Int
w <- [Summary]
xs ]
then
Summary -> OneChangeT Int [] Summary
forall a. a -> OneChangeT Int [] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> Summary
Summary Int
n (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ Int
w | Summary Int
_ Int
w <- [Summary]
xs ]))
else
[OneChangeT Int [] Summary] -> OneChangeT Int [] Summary
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum
[ Int -> Int -> Summary
Summary Int
n (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Summary] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Summary]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
newTree) Summary -> OneChangeT Int [] () -> OneChangeT Int [] Summary
forall a b. a -> OneChangeT Int [] b -> OneChangeT Int [] a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> OneChangeT Int [] ()
forall (m :: * -> *) c. Monad m => c -> OneChangeT c m ()
change Int
newNode
| Summary Int
newNode Int
newTree <- [Summary] -> [Summary]
corrections [Summary]
xs ]
corrections ::
[Summary] ->
[Summary]
corrections :: [Summary] -> [Summary]
corrections [Summary]
xs =
[ Int -> Int -> Summary
Summary (Int
nodeWeight Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
discrepency) Int
other
| (Summary Int
nodeWeight Int
treeWeight, [Summary]
xs') <- [Summary] -> [(Summary, [Summary])]
forall a. [a] -> [(a, [a])]
pickOne [Summary]
xs
, let weights :: [Int]
weights = [ Int
w | Summary Int
_ Int
w <- [Summary]
xs' ]
, [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => t a -> Bool
same [Int]
weights
, Int
other <- Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
1 [Int]
weights
, let discrepency :: Int
discrepency = Int
other Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
treeWeight
, Int
discrepency Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
]
computeCorrection :: Fix Node -> Maybe Int
computeCorrection :: Fix Node -> Maybe Int
computeCorrection = ((Maybe Int, Summary) -> Maybe Int)
-> Maybe (Maybe Int, Summary) -> Maybe Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
(=<<) (Maybe Int, Summary) -> Maybe Int
forall a b. (a, b) -> a
fst (Maybe (Maybe Int, Summary) -> Maybe Int)
-> (Fix Node -> Maybe (Maybe Int, Summary))
-> Fix Node
-> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Maybe Int, Summary)] -> Maybe (Maybe Int, Summary)
forall a. [a] -> Maybe a
listToMaybe ([(Maybe Int, Summary)] -> Maybe (Maybe Int, Summary))
-> (Fix Node -> [(Maybe Int, Summary)])
-> Fix Node
-> Maybe (Maybe Int, Summary)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OneChangeT Int [] Summary -> [(Maybe Int, Summary)]
forall c (m :: * -> *) a. OneChangeT c m a -> m (Maybe c, a)
runOneChangeT (OneChangeT Int [] Summary -> [(Maybe Int, Summary)])
-> (Fix Node -> OneChangeT Int [] Summary)
-> Fix Node
-> [(Maybe Int, Summary)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Fix Node -> OneChangeT Int [] Summary
summarize
newtype OneChangeT c m a = OCT { forall c (m :: * -> *) a. OneChangeT c m a -> m (Maybe c, a)
runOneChangeT :: m (Maybe c, a) }
change :: Monad m => c -> OneChangeT c m ()
change :: forall (m :: * -> *) c. Monad m => c -> OneChangeT c m ()
change c
c = m (Maybe c, ()) -> OneChangeT c m ()
forall c (m :: * -> *) a. m (Maybe c, a) -> OneChangeT c m a
OCT ((Maybe c, ()) -> m (Maybe c, ())
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (c -> Maybe c
forall a. a -> Maybe a
Just c
c, ()))
instance MonadPlus m => Functor (OneChangeT c m) where
fmap :: forall a b. (a -> b) -> OneChangeT c m a -> OneChangeT c m b
fmap = (a -> b) -> OneChangeT c m a -> OneChangeT c m b
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM
instance MonadPlus m => Applicative (OneChangeT c m) where
pure :: forall a. a -> OneChangeT c m a
pure a
x = m (Maybe c, a) -> OneChangeT c m a
forall c (m :: * -> *) a. m (Maybe c, a) -> OneChangeT c m a
OCT ((Maybe c, a) -> m (Maybe c, a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe c
forall a. Maybe a
Nothing, a
x))
<*> :: forall a b.
OneChangeT c m (a -> b) -> OneChangeT c m a -> OneChangeT c m b
(<*>) = OneChangeT c m (a -> b) -> OneChangeT c m a -> OneChangeT c m b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap
instance MonadPlus m => Alternative (OneChangeT c m) where
OCT m (Maybe c, a)
xs <|> :: forall a. OneChangeT c m a -> OneChangeT c m a -> OneChangeT c m a
<|> OCT m (Maybe c, a)
ys = m (Maybe c, a) -> OneChangeT c m a
forall c (m :: * -> *) a. m (Maybe c, a) -> OneChangeT c m a
OCT (m (Maybe c, a)
xs m (Maybe c, a) -> m (Maybe c, a) -> m (Maybe c, a)
forall a. m a -> m a -> m a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> m (Maybe c, a)
ys)
empty :: forall a. OneChangeT c m a
empty = m (Maybe c, a) -> OneChangeT c m a
forall c (m :: * -> *) a. m (Maybe c, a) -> OneChangeT c m a
OCT m (Maybe c, a)
forall a. m a
forall (f :: * -> *) a. Alternative f => f a
empty
instance MonadPlus m => MonadPlus (OneChangeT c m)
instance MonadPlus m => Monad (OneChangeT c m) where
OneChangeT c m a
m >>= :: forall a b.
OneChangeT c m a -> (a -> OneChangeT c m b) -> OneChangeT c m b
>>= a -> OneChangeT c m b
f = m (Maybe c, b) -> OneChangeT c m b
forall c (m :: * -> *) a. m (Maybe c, a) -> OneChangeT c m a
OCT (m (Maybe c, b) -> OneChangeT c m b)
-> m (Maybe c, b) -> OneChangeT c m b
forall a b. (a -> b) -> a -> b
$
do (Maybe c
mb1,a
x) <- OneChangeT c m a -> m (Maybe c, a)
forall c (m :: * -> *) a. OneChangeT c m a -> m (Maybe c, a)
runOneChangeT OneChangeT c m a
m
(Maybe c
mb2,b
y) <- OneChangeT c m b -> m (Maybe c, b)
forall c (m :: * -> *) a. OneChangeT c m a -> m (Maybe c, a)
runOneChangeT (a -> OneChangeT c m b
f a
x)
case (Maybe c
mb1,Maybe c
mb2) of
(Maybe c
Nothing,Maybe c
_) -> (Maybe c, b) -> m (Maybe c, b)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe c
mb2,b
y)
(Maybe c
_,Maybe c
Nothing) -> (Maybe c, b) -> m (Maybe c, b)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe c
mb1,b
y)
(Maybe c, Maybe c)
_ -> m (Maybe c, b)
forall a. m a
forall (f :: * -> *) a. Alternative f => f a
empty
data Balance
= Balanced (Set Int)
| Unbalanced (Map Int (Set Int))
deriving Int -> Balance -> ShowS
[Balance] -> ShowS
Balance -> [Char]
(Int -> Balance -> ShowS)
-> (Balance -> [Char]) -> ([Balance] -> ShowS) -> Show Balance
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Balance -> ShowS
showsPrec :: Int -> Balance -> ShowS
$cshow :: Balance -> [Char]
show :: Balance -> [Char]
$cshowList :: [Balance] -> ShowS
showList :: [Balance] -> ShowS
Show