summaryrefslogtreecommitdiff
path: root/2019/18.hs
blob: c7584465892d4538c13580f0cbbc3efff1182c14 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
{-# LANGUAGE TupleSections #-}
module Main where

import Control.Monad
import qualified Data.Array.Unboxed as A
import qualified Data.Array.ST as STA
import Data.Char
import Data.List
import qualified Data.IntMap as IntMap
import Data.IntMap (IntMap)
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Data.Maybe
import qualified Data.Set as Set
import Data.Set (Set)

-- import Debug.Trace

import Input
import qualified SmallIntSet as SIS
import SmallIntSet (SmallIntSet)


replaceAtIndex :: Int -> a -> [a] -> [a]
replaceAtIndex i _ [] = error ("Index is " ++ show i ++ " items past end in replaceAtIndex")
replaceAtIndex 0 val (_:xs) = val : xs
replaceAtIndex i val (x:xs) = x : replaceAtIndex (i-1) val xs

-- Considers a distance of '-1' to mean 'unconnected'.
-- Applies Floyd-Warshall.
transitiveClosure :: [Int] -> A.UArray (Int, Int) Int -> A.UArray (Int, Int) Int
transitiveClosure nodeList initMatrix =
    STA.runSTUArray $ do
        arr <- STA.thaw initMatrix
        forM_ nodeList $ \k ->
            forM_ nodeList $ \i ->
                forM_ nodeList $ \j -> do
                    dij <- STA.readArray arr (i, j)
                    dik <- STA.readArray arr (i, k)
                    dkj <- STA.readArray arr (k, j)
                    if dik /= -1 && dkj /= -1 && (dij == -1 || dik + dkj < dij)
                        then STA.writeArray arr (i, j) (dik + dkj)
                        else return ()
        return arr

type Pos = (Int, Int)
type Dir = (Int, Int)

reachableFrom :: Map Pos Char -> Pos -> Map Pos Int
reachableFrom bd startPos = go 0 (Set.singleton startPos) (Set.singleton startPos) Map.empty
  where
    go dist seen boundary result =
        let boundary' = [pos
                        | (x, y) <- Set.toList boundary
                        , (dx, dy) <- [(-1,0), (0,-1), (1,0), (0,1)]
                        , let pos = (x + dx, y + dy)
                        , bd Map.! pos /= '#'
                        , pos `Set.notMember` seen]
            (things, frees) = partition (\pos -> bd Map.! pos /= '.') boundary'
            result' = result <> Map.fromList (map (,dist+1) things)
        in if null frees
               then result'
               else go (dist + 1) (seen <> Set.fromList boundary') (Set.fromList frees) result'

-- [0  25   26 51  52  ]
-- [a...z,  A...Z, @...]
data Implicit = Implicit Int                        -- number of starting positions
                         (A.Array Int [Int])        -- edge list
                         (A.UArray (Int, Int) Int)  -- distance matrix, with closure taken
  deriving (Show)

codeIsStart, codeIsLower, codeIsUpper :: Int -> Bool
codeIsStart n = n >= 2 * 26
codeIsLower n = n < 26
codeIsUpper n = 26 <= n && n < 2 * 26

codeToLower :: Int -> Int
codeToLower n = n - 26     -- assumes codeIsUpper

implicitGraph :: Map Pos Char -> [Pos] -> Implicit
implicitGraph bd startPositions =
    let nStart = length startPositions
        nNodes = 2 * 26 + nStart
        posGraph = fst (goMultiple startPositions Map.empty Set.empty)
        mapGraph = Map.mapKeys (bd Map.!) (Map.map (Map.mapKeys (bd Map.!)) posGraph)
        charToNode '@' = charToNode '1'
        charToNode c | isLower c = ord c - ord 'a'
                     | isUpper c = 26 + ord c - ord 'A'
                     | isDigit c = 26 + 26 + ord c - ord '1'
                     | otherwise = undefined
        arrGraph = A.accumArray (const id) [] (0, nNodes - 1)
                        [(charToNode from, map charToNode (Map.keys tomap))
                        | (from, tomap) <- Map.assocs mapGraph]
        distArr = A.accumArray (const id) (-1) ((0, 0), (nNodes - 1, nNodes - 1))
                        [((charToNode from, charToNode to), dist)
                        | (from, tomap) <- Map.assocs mapGraph
                        , (to, dist) <- Map.assocs tomap]
        nodeList = map charToNode (Map.keys mapGraph)
    in Implicit nStart arrGraph (transitiveClosure nodeList distArr)
  where
    goMultiple :: [Pos] -> Map Pos (Map Pos Int) -> Set Pos -> (Map Pos (Map Pos Int), Set Pos)
    goMultiple curPoses graph seen = foldl' (\(gr, sn) node -> go node gr sn) (graph, seen) curPoses

    go :: Pos -> Map Pos (Map Pos Int) -> Set Pos -> (Map Pos (Map Pos Int), Set Pos)
    go curPos graph seen
        | curPos `Set.member` seen = (graph, seen)
        | otherwise =
            let reach = reachableFrom bd curPos
                newNodes = Map.keysSet reach Set.\\ seen
                graph' = Map.insert curPos reach graph
                seen' = Set.insert curPos seen
            in goMultiple (Set.toList newNodes) graph' seen'

reachable :: Implicit -> SmallIntSet -> Int -> IntMap Int
reachable (Implicit _ graph distarr) keys start = snd (go 0 (SIS.singleton start) start IntMap.empty)
  where
    go dist seen at result =
        let nexts = filter (\c -> c `SIS.notMember` seen && (codeIsStart c || codeIsLower c || codeToLower c `SIS.member` keys))
                           (graph A.! at)
            (nextPearls, nextNonpearls) = partition (\c -> codeIsLower c && c `SIS.notMember` keys) nexts
            result' = result <> IntMap.fromList [(c, dist + distarr A.! (at, c)) | c <- nextPearls]
            seen' = seen <> SIS.fromList nexts
        in -- trace ("reachable-go at=" ++ show at ++ " dist=" ++ show dist ++ " nexts=" ++ show nexts ++ " (allnexts " ++ show (graph A.! at) ++ ")") $
           if null nexts
               then (seen, result')
               else foldl (\(sn, rs) c -> go (dist + distarr A.! (at, c)) sn c rs) (seen', result') nextNonpearls

searchBFS :: Implicit -> (Int, [Int])
searchBFS implicit@(Implicit nstart _ _) =
    let startNodes = [52 + i | i <- [0..nstart-1]]
    in go 0 (Set.singleton (0, startNodes, SIS.empty, [])) Map.empty
  where
    -- pqueue: f-val, distance, nodes, keys, key order
    -- visited: nodes, keys => distance
    go :: Int -> Set (Int, [Int], SmallIntSet, [Int]) -> Map ([Int], SmallIntSet) Int -> (Int, [Int])
    go ctr pqueue visited =
        let ((dist, curnodes, keys, keyorder), newpqueue) = Set.deleteFindMin pqueue
            reachLists = [reachable implicit keys node | node <- curnodes]
            nextStates = [(dist + stepDist, stepNodes, stepKeys, stepC : keyorder)
                         | (robotIdx, reach) <- zip [0..] reachLists
                         , (stepC, stepDist) <- IntMap.assocs reach
                         , let stepKeys = SIS.insert stepC keys
                               stepNodes = replaceAtIndex robotIdx stepC curnodes
                           -- check that this next state is actually better than we've seen before
                         , maybe True (dist + stepDist <) (Map.lookup (stepNodes, stepKeys) visited)]
            visited' = Map.insert (curnodes, keys) dist visited
            pqueue' = newpqueue <> Set.fromList nextStates
            result = 
                if maybe True (dist <) (Map.lookup (curnodes, keys) visited)
                    then if all IntMap.null reachLists
                        then (dist, keyorder)
                        else go (ctr + 1) pqueue' visited'
                    else go (ctr + 1) newpqueue visited
        in -- (if ctr `rem` 20000 == 0 || all IntMap.null reachLists
           --      then trace ("go ctr=" ++ show ctr ++ " #pqueue=" ++ show (Set.size pqueue) ++ " #visited=" ++ show (Map.size visited)
           --                  ++ "  curnodes=" ++ show curnodes ++ " dist=" ++ show dist ++ " keys=" ++ show keyorder
           --                  {- ++ "  next->" ++ show nextStates -})
           --      else id)
           result

main :: IO ()
main = do
    stringbd <- getInput 18
    let bd = Map.fromList [((x, y), c) | (y, row) <- zip [0..] stringbd, (x, c) <- zip [0..] row]
        startpos = fromJust (lookup '@' (map (\(x,y) -> (y,x)) (Map.assocs bd)))

    let imgraph = implicitGraph bd [startpos]
    print (fst (searchBFS imgraph))

    let (sx, sy) = startpos
        bd2 = Map.unionWith (const id) bd
                            (Map.fromList [((sx-1,sy-1), '1'), ((sx,sy-1), '#'), ((sx+1,sy-1), '2')
                                          ,((sx-1,sy  ), '#'), ((sx,sy  ), '#'), ((sx+1,sy  ), '#')
                                          ,((sx-1,sy+1), '3'), ((sx,sy+1), '#'), ((sx+1,sy+1), '4')])
        imgraph2 = implicitGraph bd2 [(sx-1,sy-1), (sx+1,sy-1), (sx-1,sy+1), (sx+1,sy+1)]

    print (fst (searchBFS imgraph2))