Day25.hs

module Javran.AdventOfCode.Y2023.Day25 () where

import Control.Monad.State.Strict
import Control.Monad.Writer.CPS
import Data.Char
import qualified Data.IntSet as IS
import qualified Data.Map.Monoidal.Strict as MM
import qualified Data.Map.Strict as M
import qualified Data.PSQueue as PQ
import Data.Semigroup
import qualified Data.Set as S
import Javran.AdventOfCode.Misc (internalize)
import Javran.AdventOfCode.Prelude
import Text.ParserCombinators.ReadP hiding (count, get, many)

data Day25 deriving (Generic)

{-
  [General notes]

  The key here is to recognize that we are dealing with an undirected graph
  and we need to find a cut [1] for it.

  Some simple search leads to "Stoer-Wagner algorithm" [2] which is new to me
  but implementation is somewhat straightforward following along the paper [3].

  Note that this algorithm finds minimum cut for a weighted graph
  but the problem is asking for a cut of exactly 3 edges, so we need some "translation"
  in order to make this algorithm work for us:

  - Our graph is not weighted, so if we assign weight 1 to all edges,
    we are looking for a cut of total weight of 3.

  - It should be reasonable to assume that the minimum cut of input is at most 3.
    This is because if we need to cut more edges, this problem would not be solvable.

    Here we actually go further to assume the cut is exactly of weight 3, which
    seems to be the case for all inputs known to me.

  1: https://en.wikipedia.org/wiki/Cut_(graph_theory)
  2: https://en.wikipedia.org/wiki/Stoer%E2%80%93Wagner_algorithm
  3: https://scholar.google.com/scholar?cluster=10111487970680388034

 -}

{-
  Graph related types:

  - V: a vertex.

    Note that we are using a Set so to make it easier
    when it comes to merging vertices.

  - G: indicates how vertices are connected - there are redundancy in this
    as for this undirected graph (u,v) and (v,u) are both in this Map.

  - W: undirected weights assigned to edges. for keys (u,v) MinMax enforces that u < v.

 -}
type V = IS.IntSet
type G = M.Map V (S.Set V)
type W = M.Map (MinMax V) Int
type GW = (G, W)

{-

  `merge s t gw` merges vertices `s` and `t` together in graph `gw`.

  - Value of this new vertex is simply set union of `s` and `t`.

  - Edge weight between `s` and `t` is ignored (if any).

  - Connectivity outside of `s` and `t` are preserved. If they both connect to the same vertex,
    new edge weight is the sum of old weights.

 -}
merge :: V -> V -> GW -> GW
merge s t = execState do
  (g0, w0) <- get
  let
    st = s <> t
    sConns = fromMaybe S.empty (g0 M.!? s)
    tConns = fromMaybe S.empty (g0 M.!? t)
    stConns = sConns <> tConns

  -- remove s and t
  modify $ first $ M.delete t . M.delete s
  -- replace s or t with st
  modify $ first $ \g -> foldr (M.adjust (S.insert st . S.delete t . S.delete s)) g stConns
  -- move connections over to st
  modify $ first $ M.insert st (S.delete t . S.delete s $ stConns)

  -- remove edges whose one end is s or t
  modify $ second $ M.filterWithKey (\(MinMax (a, b)) _v -> notElem a [s, t] && notElem b [s, t])
  let
    stWeights :: M.Map (MinMax V) Int
    stWeights = M.fromListWith (+) do
      -- all connections except that between s and t.
      (x, conns) <- [(s, S.delete t sConns), (t, S.delete s tConns)]
      y <- S.toList conns
      Just v <- [w0 M.!? minMaxFromPair (x, y)]
      pure (minMaxFromPair (st, y), v)
  -- add new weights related to st.
  modify $ second $ M.unionWith (+) stWeights

{-
  The "MinimumCutPhase" algorithm as described in the paper.

  This feels familiar to Dijkstra's algorithm: we begin with a singleton
  set of vertices and expand it to the entire graph.

  Meanwhile a `PSQ` is used as heap to extract maximum value (with `Down`) efficiently
  to determine which vertex should be added next.

 -}
minimumCutPhase :: G -> W -> V -> Maybe ((V, V), Int)
minimumCutPhase g w a =
  fix
    ( \go aSet q acc -> case PQ.minView q of
        Nothing -> case acc of
          (s, v) : (t, _) : _ ->
            {-
              As this list is accumulated in reverse order,
              we extract first 2 elements which should be vertices we need to merge.
             -}
            Just ((s, t), v)
          _ -> Nothing
        Just (z PQ.:-> (Down zW), q1) ->
          -- insert z into `aSet`
          let
            vExtras = do
              -- collect weights related to z to update the queue.
              v <-
                maybe
                  []
                  (S.toList . (\s -> S.difference s aSet))
                  (g M.!? z)
              pure (v, w M.! minMaxFromPair (z, v))
            q2 = foldr upd q1 vExtras
              where
                upd (v, wv) =
                  PQ.alter
                    ( \case
                        Nothing -> Just (Down wv)
                        Just (Down wOld) -> Just (Down (wOld + wv))
                    )
                    v
           in
            go (S.insert z aSet) q2 ((z, zW) : acc)
    )
    (S.singleton a)
    (PQ.singleton a (Down 0))
    []

{-
  "MinimumCut" algorithm, see paper.
 -}
stoerWagner :: GW -> Maybe (ArgMin Int V)
stoerWagner p0@(g0, _) = execWriter $ runW p0
  where
    a : _ = M.keys g0
    runW = fix \go cur@(curG, curW) ->
      case minimumCutPhase curG curW a of
        Just ((s, t), w) -> do
          -- announce this cut to update current best.
          tell (Just (Min (Arg w s)))
          go (merge s t cur)
        Nothing -> pure ()

inputLineP :: ReadP (String, [String])
inputLineP = do
  let nodeP = munch1 isAsciiLower
  l <- nodeP <* strP ": "
  rs <- sepBy1 nodeP (charP ' ')
  pure (l, rs)

instance Solution Day25 where
  solutionRun _ SolutionContext {getInputS, answerShow} = do
    rawConns <- fmap (consumeOrDie inputLineP) . lines <$> getInputS
    let
      ns = S.fromList do
        (l, rs) <- rawConns
        l : rs
      (nodeToRef, _refToNode) = internalize (S.toAscList ns)
      (MM.MonoidalMap g0, w0) = mconcat do
        (l, rs) <- rawConns
        let l' = IS.singleton (nodeToRef l)
        r <- rs
        let r' = IS.singleton (nodeToRef r)
        pure
          ( MM.fromList
              [ (l', S.singleton r')
              , (r', S.singleton l')
              ]
          , M.singleton (minMaxFromPair (l', r')) 1
          )

    let
      Just (Min (Arg 3 t)) = stoerWagner (g0, w0)
      part = IS.size t
    answerShow (part * (M.size g0 - part))