I would like to compute the 8-way neighbors of every element in a rectangular list of lists xss :: [[a]], preserving the original element and defaulting on the boundary. In particular, I want to write a function neighbors :: a -> [[a]] -> [[NonEmpty a]] where each inner list has length 9, with the head being the original element.
Here is the output I'm expecting:
ghci> mapM_ print [[1,2,3],[4,5,6],[7,8,9]]
[1,2,3]
[4,5,6]
[7,8,9]
ghci> mapM_ (mapM_ print) $ neighbors 0 [[1,2,3],[4,5,6],[7,8,9]]
[1,0,0,0,0,2,0,4,5]
[2,0,0,0,1,3,4,5,6]
[3,0,0,0,2,0,5,6,0]
[4,0,1,2,0,5,0,7,8]
[5,1,2,3,4,6,7,8,9]
[6,2,3,0,5,0,8,9,0]
[7,0,4,5,0,8,0,0,0]
[8,4,5,6,7,9,0,0,0]
[9,5,6,0,8,0,0,0,0]
I have read the nice answers to this old question, but it doesn't quite satisfy these requirements because it generates pairs and destroys the structure of the original matrix.
However I am hoping it is possible to write something similarly concise.
I've come up with two very bad partial solutions, one which performs a ton of rotations and another which is not coordinate-free:
-- | Add a ring of default elements around a 2d array.
augment :: a -> [[a]] -> [[a]]
augment d m = map (pad d) . pad (replicate ncols d) $ m
where
pad d' = reverse . (d' :) . reverse . (d' :)
ncols = foldr (max . length) 0 m
-- | Neighbors in a 2d array, with default.
neighbors1 :: a -> [[a]] -> [[[a]]]
neighbors1 d m = map transpose $ transpose $ map ($ augment d m) [middle, upL, up, upR, left, right, downL, down, downR]
where
rot = transpose . reverse
right = map (drop 2) . both (drop 1)
up = rot . rot . rot . right . rot
left = rot . rot . right . rot . rot
down = rot . right . rot . rot . rot
downR = map (drop 2) . drop 2
upR = rot . rot . rot . downR . rot
upL = rot . rot . downR . rot . rot
downL = rot . downR . rot . rot . rot
middle = map (both (drop 1)) . both (drop 1)
both f = reverse . f . reverse . f
-- | Includes center element as head.
-- O(mn*log(mn)) Data.Map.Strict
-- O(mn) Data.HashMap.Strict
neighbors :: a -> [[a]] -> [[[a]]]
neighbors d zss = unmaybe [[map (vals !?) (nbs x y) | y <- [0..n-1]] | x <- [0..m-1]]
where
unmaybe = (map . map . map) (fromMaybe d)
(m, n) = (length zss, maybe 0 length $ listToMaybe zss)
nbs x y = (x,y) : [(i,j) | i <- [x-1..x+1], j <- [y-1..y+1], (i,j) /= (x,y)]
vals = M.fromAscList [((i,j),z) | (i,zs) <- zip [0..] zss, (j,z) <- zip [0..] zs]
Also note that the augment function is a bit gross. Perhaps it can be done more elegantly with zipWith const? But I can't figure out how. I also tried using Maybes instead of defaults, which would also be acceptable, but once I have an array of the form [Maybe [Maybe a]], I can't figure out how to operate on it.
Edit: I've come up with a slightly better version that I will include for posterity, but I think the answers so far are probably still superior. I still couldn't figure out the NonEmpty bit.
This still makes use of this ugly augment function. I wrote another version of it but they both seem inelegant.
rev :: ([a] -> [a]) -> [a] -> [a]
rev f = reverse . f . reverse
both :: ([a] -> [a]) -> [a] -> [a]
both f = rev f . f
-- | Add a ring of default elements around a 2d array.
augment :: a -> [[a]] -> [[a]]
augment d xss = map (pad d) . pad (replicate n d) $ xss
where
pad d' = both (d':)
n = maybe 0 length (listToMaybe xss)
-- | Add a ring of default elements around a 2d array.
augment' :: a -> [[a]] -> [[a]]
augment' _ [] = []
augment' d xss@(xs:_) = map (pad d) . pad ds $ xss
where
pad d' = both (d':)
ds = map (const d) xs
-- | Neighbors in a 2d array, with default.
neighbors1 :: a -> [[a]] -> [[[a]]]
neighbors1 d m = map transpose $ transpose $ map ($ augment d m)
-- vertical horizontal
[ center . map center
, center . map down
, center . map up
, down . map center
, up . map center
, down . map down
, down . map up
, up . map down
, up . map up
]
where
center = both (drop 1)
down = drop 2
up = rev down
Here is an adapted version of Li-yao Xia's solution. I really like the idea of arranging the transforms in the list in the shape of a 3x3 grid, but I wanted to make it easier to textually transform this back and forth from 4-way neighbors to 8-way neighbors, and from the NonEmpty version to the one that excludes the middle.
-- Rotate the outermost axis to the innermost position.
transpose3D :: [[[a]]] -> [[[a]]]
transpose3D = map transpose . transpose
zipWith2D :: (a -> b -> c) -> [[a]] -> [[b]] -> [[c]]
zipWith2D = zipWith . zipWith
neighbors6 :: a -> [[a]] -> [[NonEmpty a]]
neighbors6 d m = zipWith2D (:|) m $ transpose3D $ map ($ m)
[ up
, down
, left
, right
, up . left
, up . right
, down . left
, down . right
]
where
up = (repeat d :)
down = (++ [repeat d]) . drop 1
left = map (d :)
right = map ((++ [d]) . drop 1)
In the above we can just delete the zipWith2D call and it becomes [[[a]]] with only neighbors, no center element, and we can also just delete the last four lines of the list and get cardinals.
It's worth noting that his answer is faster than anything else I've tried in a very rudimentary benchmark using deepseq.