{-# Language ImportQualifiedPost, BlockArguments #-}
{-|
Module      : Advent.Prelude
Description : Prelude extension for AoC solutions
Copyright   : (c) Eric Mertens, 2018-2021
License     : ISC
Maintainer  : emertens@gmail.com

Various helper functions for common operations needed in
Advent of Code problems.

-}
module Advent.Prelude where

import Control.Monad.Trans.State (StateT(StateT, runStateT))
import Data.Array.Unboxed qualified as A
import Data.Coerce (coerce)
import Data.Foldable (toList)
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap
import Data.List (foldl', inits, sortBy, tails, mapAccumL)
import Data.Map (Map)
import Data.Map.Strict qualified as SMap
import Data.MemoTrie (HasTrie, memo3, mup)
import Data.Ord (comparing)
import Data.Set (Set)
import Data.Set qualified as Set

-- | Count the number of elements in a foldable value that satisfy a predicate.
count :: (Foldable f, Eq a) => a -> f a -> Int
count :: forall (f :: * -> *) a. (Foldable f, Eq a) => a -> f a -> Int
count a
x = (a -> Bool) -> f a -> Int
forall (f :: * -> *) a. Foldable f => (a -> Bool) -> f a -> Int
countBy (a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==)

-- | Count the number of elements in a foldable value that satisfy a predicate.
countBy :: Foldable f => (a -> Bool) -> f a -> Int
countBy :: forall (f :: * -> *) a. Foldable f => (a -> Bool) -> f a -> Int
countBy a -> Bool
p = (Int -> a -> Int) -> Int -> f a -> Int
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
acc a
x -> if a -> Bool
p a
x then Int
accInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1 else Int
acc) Int
0

-- | Return true when the whole list is comprised of equal elements.
--
-- >>> same [1,1,1]
-- True
-- >>> same []
-- True
-- >>> same [1]
-- True
-- >>> same [1,1,2]
-- False
same :: Foldable t => Eq a => t a -> Bool
same :: forall (t :: * -> *) a. (Foldable t, Eq a) => t a -> Bool
same t a
xs = (a -> Bool) -> t a -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ([a] -> a
forall a. [a] -> a
head (t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t a
xs) a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==) t a
xs

-- | Returns a list of ways to select an element from a list without
-- replacement.
--
-- >>> pickOne []
-- []
-- >>> pickOne [1]
-- [(1,[])]
-- >>> pickOne [1,2,3]
-- [(1,[2,3]),(2,[1,3]),(3,[1,2])]
pickOne :: [a] -> [(a, [a])]
pickOne :: forall a. [a] -> [(a, [a])]
pickOne [a]
xs = [ (a
x, [a]
l[a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++[a]
r) | ([a]
l,a
x:[a]
r) <- [[a]] -> [[a]] -> [([a], [a])]
forall a b. [a] -> [b] -> [(a, b)]
zip ([a] -> [[a]]
forall a. [a] -> [[a]]
inits [a]
xs) ([a] -> [[a]]
forall a. [a] -> [[a]]
tails [a]
xs) ]

-- | Implementation of 'Data.List.nub' that uses 'Ord' for efficiency.
ordNub :: Ord a => [a] -> [a]
ordNub :: forall a. Ord a => [a] -> [a]
ordNub [a]
xs = (a -> (Set a -> [a]) -> Set a -> [a])
-> (Set a -> [a]) -> [a] -> Set a -> [a]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> (Set a -> [a]) -> Set a -> [a]
forall {a}. Ord a => a -> (Set a -> [a]) -> Set a -> [a]
f ([a] -> Set a -> [a]
forall a b. a -> b -> a
const []) [a]
xs Set a
forall a. Set a
Set.empty
  where
    f :: a -> (Set a -> [a]) -> Set a -> [a]
f a
x Set a -> [a]
rec Set a
seen =
      case Set a -> [a]
rec (Set a -> [a]) -> (Bool, Set a) -> (Bool, [a])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Bool -> (Bool, Bool)) -> a -> Set a -> (Bool, Set a)
forall a (f :: * -> *).
(Ord a, Functor f) =>
(Bool -> f Bool) -> a -> Set a -> f (Set a)
Set.alterF (\Bool
old -> (Bool
old, Bool
True)) a
x Set a
seen of
        (Bool
True,  [a]
ys) -> [a]
ys
        (Bool
False, [a]
ys) -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ys

-- | Compute the minimum element of a list or return Nothing if it is empty.
--
-- >>> minimumMaybe []
-- Nothing
-- >>> minimumMaybe [2,1,3]
-- Just 1
minimumMaybe :: Ord a => [a] -> Maybe a
minimumMaybe :: forall a. Ord a => [a] -> Maybe a
minimumMaybe [a]
xs
  | [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs   = Maybe a
forall a. Maybe a
Nothing
  | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [a]
xs


-- | Compute the maximum element of a list or return Nothing if it is empty.
--
-- >>> maximumMaybe []
-- Nothing
-- >>> maximumMaybe [2,1,3]
-- Just 3
maximumMaybe :: Ord a => [a] -> Maybe a
maximumMaybe :: forall a. Ord a => [a] -> Maybe a
maximumMaybe [a]
xs
  | [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [a]
xs   = Maybe a
forall a. Maybe a
Nothing
  | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [a]
xs

-- | Compute the number of occurrences of the elements in a given list.
--
-- >>> counts "bababc"
-- fromList [('a',2),('b',3),('c',1)]
counts :: (Foldable f, Ord a) => f a -> Map a Int
counts :: forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts f a
xs = (Int -> Int -> Int) -> [(a, Int)] -> Map a Int
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
SMap.fromListWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [(a
x,Int
1) | a
x <- f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f a
xs]

-- | Compose a list of functions together
--
-- >>> compose [ (1:), (2:), (3:) ] []
-- [1,2,3]
compose :: [a -> a] -> a -> a
compose :: forall a. [a -> a] -> a -> a
compose = ((a -> a) -> (a -> a) -> a -> a) -> (a -> a) -> [a -> a] -> a -> a
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) a -> a
forall a. a -> a
id

-- | Split list into chunks. The last chunk might be incomplete.
--
-- >>> chunks 3 [1..9]
-- [[1,2,3],[4,5,6],[7,8,9]]
--
-- >>> chunks 3 [1..7]
-- [[1,2,3],[4,5,6],[7]]
--
-- >>> chunks 3 []
-- []
chunks :: Int -> [a] -> [[a]]
chunks :: forall a. Int -> [a] -> [[a]]
chunks Int
_ [] = []
chunks Int
n [a]
xs =
  case Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
xs of
    ([a]
a,[a]
b) -> [a]
a [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: Int -> [a] -> [[a]]
forall a. Int -> [a] -> [[a]]
chunks Int
n [a]
b

-- | Löb's theorem
--
-- <https://github.com/quchen/articles/blob/master/loeb-moeb.md>
-- <https://en.wikipedia.org/wiki/L%C3%B6b%27s_theorem>
löb :: Functor f => f (f a -> a) -> f a
löb :: forall (f :: * -> *) a. Functor f => f (f a -> a) -> f a
löb = (((f a -> a) -> a) -> f (f a -> a) -> f a) -> f (f a -> a) -> f a
forall a b c. (((a -> b) -> b) -> c -> a) -> c -> a
möb ((f a -> a) -> a) -> f (f a -> a) -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap

-- | 'löb' generalized over 'fmap'
möb :: (((a -> b) -> b) -> c -> a) -> c -> a
möb :: forall a b c. (((a -> b) -> b) -> c -> a) -> c -> a
möb ((a -> b) -> b) -> c -> a
f = \c
x -> let go :: a
go = ((a -> b) -> b) -> c -> a
f ((a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
$ a
go) c
x in a
go

-- | Index an array returning 'Nothing' if the index is out of bounds.
arrIx :: (A.IArray a e, A.Ix i) => a i e -> i -> Maybe e
arrIx :: forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> Maybe e
arrIx a i e
a i
i
  | (i, i) -> i -> Bool
forall a. Ix a => (a, a) -> a -> Bool
A.inRange (a i e -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
A.bounds a i e
a) i
i = e -> Maybe e
forall a. a -> Maybe a
Just (e -> Maybe e) -> e -> Maybe e
forall a b. (a -> b) -> a -> b
$! a i e
a a i e -> i -> e
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
A.! i
i
  | Bool
otherwise = Maybe e
forall a. Maybe a
Nothing

-- | Apply a function @n@ times strictly.
times :: Int -> (a -> a) -> a -> a
times :: forall a. Int -> (a -> a) -> a -> a
times Int
n a -> a
f a
x
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = a
x
  | Bool
otherwise = Int -> (a -> a) -> a -> a
forall a. Int -> (a -> a) -> a -> a
times (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) a -> a
f (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$! a -> a
f a
x

-- | Given a list of constraints such that each constraint identifies
-- a unique variable and the set of assignments it can have, this
-- computes assignments of those variables such that no two input
-- variables are assigned the same value.
--
-- >>> uniqueAssignment [Set.fromList "ab", Set.fromList "bc"]
-- ["ab","ac","bc"]
uniqueAssignment ::
  (Traversable t, Ord a) =>
  t (Set a) {- ^ element must map to one of the corresponding set members -} ->
  [t a]     {- ^ possible assignments -}
uniqueAssignment :: forall (t :: * -> *) a.
(Traversable t, Ord a) =>
t (Set a) -> [t a]
uniqueAssignment t (Set a)
m =
  [ ([a], t a) -> t a
forall a b. (a, b) -> b
snd (([a] -> Set a -> ([a], a)) -> [a] -> t (Set a) -> ([a], t a)
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL (\(a
x:[a]
xs) Set a
_ -> ([a]
xs,a
x)) (IntMap a -> [a]
forall a. IntMap a -> [a]
IntMap.elems IntMap a
a) t (Set a)
m)
  | IntMap a
a <- IntMap a -> [(Int, Set a)] -> [IntMap a]
forall a. Ord a => IntMap a -> [(Int, Set a)] -> [IntMap a]
go IntMap a
forall a. IntMap a
IntMap.empty ([Int] -> [Set a] -> [(Int, Set a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] (t (Set a) -> [Set a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t (Set a)
m))]
  where
    go :: Ord a => IntMap a -> [(Int, Set a)] -> [IntMap a]
    go :: forall a. Ord a => IntMap a -> [(Int, Set a)] -> [IntMap a]
go IntMap a
a [(Int, Set a)]
xs =
      case ((Int, Set a) -> (Int, Set a) -> Ordering)
-> [(Int, Set a)] -> [(Int, Set a)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((Int, Set a) -> Int) -> (Int, Set a) -> (Int, Set a) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Set a -> Int
forall a. Set a -> Int
Set.size (Set a -> Int) -> ((Int, Set a) -> Set a) -> (Int, Set a) -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Set a) -> Set a
forall a b. (a, b) -> b
snd)) [(Int, Set a)]
xs of
        [] -> [IntMap a
a]
        (Int
k,Set a
vs):[(Int, Set a)]
rest ->
          do a
v <- Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
vs
             IntMap a -> [(Int, Set a)] -> [IntMap a]
forall a. Ord a => IntMap a -> [(Int, Set a)] -> [IntMap a]
go (Int -> a -> IntMap a -> IntMap a
forall a. Int -> a -> IntMap a -> IntMap a
IntMap.insert Int
k a
v IntMap a
a) ((Set a -> Set a) -> (Int, Set a) -> (Int, Set a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete a
v) ((Int, Set a) -> (Int, Set a)) -> [(Int, Set a)] -> [(Int, Set a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Int, Set a)]
rest)

-- | Convert a big-endian list of digits to a single number.
--
-- >>> fromDigits 10 [1,2,3,4]
-- 1234
--
-- >>> fromDigits 2 [12]
-- 12
--
-- >>> fromDigits 10 []
-- 0
fromDigits :: Integral a => a -> [a] -> a
fromDigits :: forall a. Integral a => a -> [a] -> a
fromDigits a
base
  | a
base a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2  = [Char] -> [a] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"fromDigits: bad base"
  | Bool
otherwise = (a -> a -> a) -> a -> [a] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\a
acc a
x -> a
acc a -> a -> a
forall a. Num a => a -> a -> a
* a
base a -> a -> a
forall a. Num a => a -> a -> a
+ a
x) a
0

-- | Convert a number to a list of digits in a given radix.
--
-- >>> toDigits 2 12
-- [1,1,0,0]
--
-- >>> toDigits 10 1234
-- [1,2,3,4]
--
-- >>> toDigits 10 0
-- []
toDigits :: Integral a => a -> a -> [a]
toDigits :: forall a. Integral a => a -> a -> [a]
toDigits a
base a
x
  | a
base a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
2  = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"toDigits: bad base"
  | a
x a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0     = [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"toDigits: negative number"
  | Bool
otherwise = [a] -> a -> [a]
go [] a
x
  where
    go :: [a] -> a -> [a]
go [a]
xs a
0 = [a]
xs
    go [a]
xs a
n = case a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
quotRem a
n a
base of
                (a
n', a
digit) -> [a] -> a -> [a]
go (a
digita -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
xs) a
n'

-- | Efficient exponentiation using an associative operator
--
-- >>> power (+) 1 10
-- 10
--
-- >>> power (*) 2 10
-- 1024 
power :: (a -> a -> a) -> a -> Integer -> a
power :: forall a. (a -> a -> a) -> a -> Integer -> a
power a -> a -> a
(#) a
one Integer
n
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
1 = [Char] -> a
forall a. HasCallStack => [Char] -> a
error ([Char]
"power: bad argument " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
n)
  | Bool
otherwise = Integer -> a
forall {a}. Integral a => a -> a
go Integer
n
  where
    double :: a -> a
double a
x = a
x a -> a -> a
# a
x
    go :: a -> a
go a
i
      | a
1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
i    = a
one
      | a -> Bool
forall a. Integral a => a -> Bool
even a
i    = a -> a
double (a -> a
go (a
i a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
2))
      | Bool
otherwise = a -> a
double (a -> a
go (a
i a -> a -> a
forall a. Integral a => a -> a -> a
`quot` a
2)) a -> a -> a
# a
one

scanlM :: (Traversable t, Monad m) => (b -> a -> m (c, a)) -> a -> t b -> m (t c, a)
scanlM :: forall (t :: * -> *) (m :: * -> *) b a c.
(Traversable t, Monad m) =>
(b -> a -> m (c, a)) -> a -> t b -> m (t c, a)
scanlM b -> a -> m (c, a)
f a
z t b
t = StateT a m (t c) -> a -> m (t c, a)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT ((b -> StateT a m c) -> t b -> StateT a m (t c)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((b -> a -> m (c, a)) -> b -> StateT a m c
coerce b -> a -> m (c, a)
f) t b
t) a
z