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
178
179
180
181
182
183
|
{-# 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 :: SmallIntSet -> Implicit -> (Int, [Int])
searchBFS allKeys implicit@(Implicit nstart _ distarr) =
let startNodes = [52 + i | i <- [0..nstart-1]]
in go 0 (Set.singleton (heuristic startNodes SIS.empty, 0, startNodes, SIS.empty, [])) Map.empty
where
-- pqueue: f-val, distance, nodes, keys, key order
-- visited: nodes, keys => distance
go :: Int -> Set (Int, Int, [Int], SmallIntSet, [Int]) -> Map ([Int], SmallIntSet) Int -> (Int, [Int])
go ctr pqueue visited =
let ((heurval, dist, curnodes, keys, keyorder), newpqueue) = Set.deleteFindMin pqueue
reachLists = [reachable implicit keys node | node <- curnodes]
nextStates = [(dist + stepDist + heuristic stepNodes stepKeys, 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 all IntMap.null reachLists
then if heurval == dist then (dist, keyorder) else error ("heurval - dist = " ++ show (heurval - dist) ++ " in terminal state!")
else go (ctr + 1) pqueue' visited'
in -- (if ctr `rem` 20000 == 0 || all IntMap.null reachLists
-- then trace ("go #pqueue=" ++ show (Set.size pqueue) ++ " #visited=" ++ show (Map.size visited)
-- ++ " curnodes=" ++ show curnodes ++ " dist=" ++ show dist ++ " heurval=" ++ show heurval ++ " keys=" ++ show keyorder
-- {- ++ " next->" ++ show nextStates -})
-- else id)
result
heuristic :: [Int] -> SmallIntSet -> Int
heuristic _curnodes keys =
let remainKeys = allKeys SIS.\\ keys
allDists = [distarr A.! (x, y) | x:xs <- tails (SIS.toList remainKeys), y <- xs]
distLowerBound = sum (take (SIS.size remainKeys - 1) (sort (filter (/= -1) allDists)))
in distLowerBound
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]
allKeys = SIS.fromList [ord c - ord 'a' | c <- Map.elems bd, isLower c]
print (fst (searchBFS allKeys 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 allKeys imgraph2))
|
