{-# OPTIONS_GHC -Wno-x-partial -Wno-unrecognised-warning-flags #-}

{- |
Description: Construct minimum spanning tree
Copyright: Copyright (C) 2021 Yoo Chung
License: GPL-3.0-or-later
Maintainer: dev@chungyc.org

Some solutions to "Problems.P84" of Ninety-Nine Haskell "Problems".
-}
module Solutions.P84 (minimumSpanningTree) where

import           Data.Map.Lazy   (Map)
import qualified Data.Map.Lazy   as Map
import           Data.Maybe      (fromJust)
import           Data.Set        (Set)
import qualified Data.Set        as Set
import           Problems.Graphs

-- | Write a function which constructs the minimum spanning tree of a given weighted graph.
-- While the weight of an edge could be encoded in the graph represention itself,
-- here we will specify the weight of each edge in a separate map.
minimumSpanningTree :: G -> Map Edge Int -> G
minimumSpanningTree :: G -> Map Edge Vertex -> G
minimumSpanningTree g :: G
g@(G Map Vertex (Set Vertex)
m) Map Edge Vertex
weights
  | Set Vertex -> Bool
forall a. Set a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (G -> Set Vertex
forall g. Graph g => g -> Set Vertex
vertexes G
g) = Map Vertex (Set Vertex) -> G
G Map Vertex (Set Vertex)
forall k a. Map k a
Map.empty
  | Bool
otherwise = Set Vertex -> Set Edge -> G -> G
toSpanningTree Set Vertex
vs Set Edge
es G
g
  where wg :: (G, Map Edge Vertex)
wg = (G
g, Map Edge Vertex
weights)
        (Vertex
v, Set Vertex
_) = Map Vertex (Set Vertex) -> (Vertex, Set Vertex)
forall k a. Map k a -> (k, a)
Map.findMin Map Vertex (Set Vertex)
m  -- chosen arbitrarily
        (Set Vertex
vs, Set Edge
es, Map Vertex [Edge]
_) = (G, Map Edge Vertex)
-> (Set Vertex, Set Edge, Map Vertex [Edge])
-> (Set Vertex, Set Edge, Map Vertex [Edge])
expand (G, Map Edge Vertex)
wg (Vertex -> Set Vertex
forall a. a -> Set a
Set.singleton Vertex
v, Set Edge
forall a. Set a
Set.empty, (G, Map Edge Vertex) -> Vertex -> Set Vertex -> Map Vertex [Edge]
weightEdgesFromVertex (G, Map Edge Vertex)
wg Vertex
v Set Vertex
forall a. Set a
Set.empty)

type WeightedG = (G, Map Edge Int)

-- | Partially constructed minimum spanning tree.
type Partial = (Set Vertex,      -- Vertexes in the partially constructed tree.
                Set Edge,        -- Edges in the partially constructed tree.
                Map Int [Edge])  -- Boundary of edges between inside the tree and outside the tree, keyed by weight.

-- | Converts the given vertexes and edges into a graph only if they from a spanning tree of the given graph.
-- Otherwise, returns an empty graph.
toSpanningTree :: Set Vertex -> Set Edge -> G -> G
toSpanningTree :: Set Vertex -> Set Edge -> G -> G
toSpanningTree Set Vertex
vs Set Edge
es G
g
  | Set Vertex
vs Set Vertex -> Set Vertex -> Bool
forall a. Eq a => a -> a -> Bool
== G -> Set Vertex
forall g. Graph g => g -> Set Vertex
vertexes G
g = Maybe G -> G
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe G -> G) -> Maybe G -> G
forall a b. (a -> b) -> a -> b
$ (Set Vertex, Set Edge) -> Maybe G
forall g. Graph g => (Set Vertex, Set Edge) -> Maybe g
toGraph (Set Vertex
vs, Set Edge
es)
  | Bool
otherwise        = Map Vertex (Set Vertex) -> G
G Map Vertex (Set Vertex)
forall k a. Map k a
Map.empty

-- | Expand the partially constructed minimum spanning tree by one edge.
expand :: WeightedG -> Partial -> Partial
expand :: (G, Map Edge Vertex)
-> (Set Vertex, Set Edge, Map Vertex [Edge])
-> (Set Vertex, Set Edge, Map Vertex [Edge])
expand (G, Map Edge Vertex)
wg r :: (Set Vertex, Set Edge, Map Vertex [Edge])
r@(Set Vertex
vs, Set Edge
_, Map Vertex [Edge]
boundary)
  | Map Vertex [Edge] -> Bool
forall k a. Map k a -> Bool
Map.null Map Vertex [Edge]
boundary = (Set Vertex, Set Edge, Map Vertex [Edge])
r
  | Bool
otherwise = (G, Map Edge Vertex)
-> Maybe (Edge, Map Vertex [Edge])
-> (Set Vertex, Set Edge, Map Vertex [Edge])
-> (Set Vertex, Set Edge, Map Vertex [Edge])
incorporateMinEdge (G, Map Edge Vertex)
wg Maybe (Edge, Map Vertex [Edge])
e (Set Vertex, Set Edge, Map Vertex [Edge])
r
  where e :: Maybe (Edge, Map Vertex [Edge])
e = (G, Map Edge Vertex)
-> Map Vertex [Edge]
-> Set Vertex
-> Maybe (Edge, Map Vertex [Edge])
extractMinEdge (G, Map Edge Vertex)
wg Map Vertex [Edge]
boundary Set Vertex
vs

-- | Extract the edge with minimum weight which connects to outside the tree from the boundary.
--
-- While edges that do not connect to the outside of the tree are not added,
-- edges that between vertexes inside the tree can exist because they are not removed.
-- Such edges are skipped.
extractMinEdge :: WeightedG -> Map Int [Edge] -> Set Vertex -> Maybe (Edge, Map Int [Edge])
extractMinEdge :: (G, Map Edge Vertex)
-> Map Vertex [Edge]
-> Set Vertex
-> Maybe (Edge, Map Vertex [Edge])
extractMinEdge (G, Map Edge Vertex)
wg Map Vertex [Edge]
boundary Set Vertex
vs
  | Map Vertex [Edge] -> Bool
forall k a. Map k a -> Bool
Map.null Map Vertex [Edge]
boundary = Maybe (Edge, Map Vertex [Edge])
forall a. Maybe a
Nothing
  | Bool
otherwise         = Edge -> Maybe (Edge, Map Vertex [Edge])
check Edge
e
  where (Edge
e, Map Vertex [Edge]
boundary') = (G, Map Edge Vertex)
-> Map Vertex [Edge] -> Set Vertex -> (Edge, Map Vertex [Edge])
extractFromBoundary (G, Map Edge Vertex)
wg Map Vertex [Edge]
boundary Set Vertex
vs
        check :: Edge -> Maybe (Edge, Map Vertex [Edge])
check (Edge (Vertex
u', Vertex
v'))
          | Vertex -> Set Vertex -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Vertex
u' Set Vertex
vs Bool -> Bool -> Bool
&& Vertex -> Set Vertex -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Vertex
v' Set Vertex
vs = (G, Map Edge Vertex)
-> Map Vertex [Edge]
-> Set Vertex
-> Maybe (Edge, Map Vertex [Edge])
extractMinEdge (G, Map Edge Vertex)
wg Map Vertex [Edge]
boundary' Set Vertex
vs
          | Bool
otherwise                            = (Edge, Map Vertex [Edge]) -> Maybe (Edge, Map Vertex [Edge])
forall a. a -> Maybe a
Just (Edge
e, Map Vertex [Edge]
boundary')

-- | Extract the edge with minimum weight from the boundary.
--
-- The edge may not connect to a vertex outside the partially constructed tree.
extractFromBoundary :: WeightedG -> Map Int [Edge] -> Set Vertex -> (Edge, Map Int [Edge])
extractFromBoundary :: (G, Map Edge Vertex)
-> Map Vertex [Edge] -> Set Vertex -> (Edge, Map Vertex [Edge])
extractFromBoundary (G, Map Edge Vertex)
wg Map Vertex [Edge]
boundary Set Vertex
vs = (Edge
e, Map Vertex [Edge]
boundary')
  where ((Vertex
minWeight, [Edge]
es), Map Vertex [Edge]
boundary'') = Map Vertex [Edge] -> ((Vertex, [Edge]), Map Vertex [Edge])
forall k a. Map k a -> ((k, a), Map k a)
Map.deleteFindMin Map Vertex [Edge]
boundary
        e :: Edge
e = [Edge] -> Edge
forall a. HasCallStack => [a] -> a
head [Edge]
es
        v :: Vertex
v = Edge -> Set Vertex -> Vertex
newVertex Edge
e Set Vertex
vs
        boundary''' :: Map Vertex [Edge]
boundary''' = (Vertex, [Edge]) -> Map Vertex [Edge] -> Map Vertex [Edge]
forall {k} {a}. Ord k => (k, [a]) -> Map k [a] -> Map k [a]
reinsert (Vertex
minWeight, [Edge] -> [Edge]
forall a. HasCallStack => [a] -> [a]
tail [Edge]
es) Map Vertex [Edge]
boundary''
        boundary' :: Map Vertex [Edge]
boundary' = ([Edge] -> [Edge] -> [Edge])
-> Map Vertex [Edge] -> Map Vertex [Edge] -> Map Vertex [Edge]
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith [Edge] -> [Edge] -> [Edge]
forall a. [a] -> [a] -> [a]
(++) Map Vertex [Edge]
edgesFromVertex Map Vertex [Edge]
boundary'''
        edgesFromVertex :: Map Vertex [Edge]
edgesFromVertex = (G, Map Edge Vertex) -> Vertex -> Set Vertex -> Map Vertex [Edge]
weightEdgesFromVertex (G, Map Edge Vertex)
wg Vertex
v Set Vertex
vs
        reinsert :: (k, [a]) -> Map k [a] -> Map k [a]
reinsert (k
_, []) Map k [a]
m'  = Map k [a]
m'
        reinsert (k
w, [a]
es') Map k [a]
m' = k -> [a] -> Map k [a] -> Map k [a]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
w [a]
es' Map k [a]
m'

-- | Incorporate an edge into a partially constructed minimum spanning tree.
incorporateMinEdge :: WeightedG -> Maybe (Edge, Map Int [Edge]) -> Partial -> Partial
incorporateMinEdge :: (G, Map Edge Vertex)
-> Maybe (Edge, Map Vertex [Edge])
-> (Set Vertex, Set Edge, Map Vertex [Edge])
-> (Set Vertex, Set Edge, Map Vertex [Edge])
incorporateMinEdge (G, Map Edge Vertex)
_ Maybe (Edge, Map Vertex [Edge])
Nothing (Set Vertex, Set Edge, Map Vertex [Edge])
r = (Set Vertex, Set Edge, Map Vertex [Edge])
r
incorporateMinEdge (G, Map Edge Vertex)
wg (Just (e :: Edge
e@(Edge (Vertex
u, Vertex
v)), Map Vertex [Edge]
boundary')) (Set Vertex
vs, Set Edge
es, Map Vertex [Edge]
_) = (G, Map Edge Vertex)
-> (Set Vertex, Set Edge, Map Vertex [Edge])
-> (Set Vertex, Set Edge, Map Vertex [Edge])
expand (G, Map Edge Vertex)
wg (Set Vertex
vs', Set Edge
es', Map Vertex [Edge]
boundary')
  where es' :: Set Edge
es' = Edge -> Set Edge -> Set Edge
forall a. Ord a => a -> Set a -> Set a
Set.insert Edge
e Set Edge
es
        vs' :: Set Vertex
vs' = Vertex -> Set Vertex -> Set Vertex
forall a. Ord a => a -> Set a -> Set a
Set.insert Vertex
v (Set Vertex -> Set Vertex) -> Set Vertex -> Set Vertex
forall a b. (a -> b) -> a -> b
$ Vertex -> Set Vertex -> Set Vertex
forall a. Ord a => a -> Set a -> Set a
Set.insert Vertex
u Set Vertex
vs

-- | Returns the weight to edges map for edges connected to a vertex.
--
-- Edges that do not connect to outside the partially constructed tree are excluded.
weightEdgesFromVertex :: WeightedG -> Vertex -> Set Vertex -> Map Int [Edge]
weightEdgesFromVertex :: (G, Map Edge Vertex) -> Vertex -> Set Vertex -> Map Vertex [Edge]
weightEdgesFromVertex (G Map Vertex (Set Vertex)
m, Map Edge Vertex
weights) Vertex
v Set Vertex
vs =
  Set Edge -> Set Vertex -> Map Edge Vertex -> Map Vertex [Edge]
weightEdges ((Vertex -> Edge) -> Set Vertex -> Set Edge
forall b a. Ord b => (a -> b) -> Set a -> Set b
Set.map (\Vertex
v' -> (Vertex, Vertex) -> Edge
Edge (Vertex
v,Vertex
v')) (Set Vertex -> Set Edge) -> Set Vertex -> Set Edge
forall a b. (a -> b) -> a -> b
$ Set Vertex -> Vertex -> Map Vertex (Set Vertex) -> Set Vertex
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Set Vertex
forall a. Set a
Set.empty Vertex
v Map Vertex (Set Vertex)
m) Set Vertex
vs Map Edge Vertex
weights

-- | Returns the weight to edges map from the edges.
--
-- Edges that do not connect to outside the partially constructed tree are excluded.
weightEdges :: Set Edge -> Set Vertex -> Map Edge Int -> Map Int [Edge]
weightEdges :: Set Edge -> Set Vertex -> Map Edge Vertex -> Map Vertex [Edge]
weightEdges Set Edge
es Set Vertex
vs Map Edge Vertex
weights = ([Edge] -> [Edge] -> [Edge])
-> [(Vertex, [Edge])] -> Map Vertex [Edge]
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith [Edge] -> [Edge] -> [Edge]
forall a. [a] -> [a] -> [a]
(++) ([(Vertex, [Edge])] -> Map Vertex [Edge])
-> [(Vertex, [Edge])] -> Map Vertex [Edge]
forall a b. (a -> b) -> a -> b
$ (Edge -> (Vertex, [Edge])) -> [Edge] -> [(Vertex, [Edge])]
forall a b. (a -> b) -> [a] -> [b]
map Edge -> (Vertex, [Edge])
weightEdge ([Edge] -> [(Vertex, [Edge])]) -> [Edge] -> [(Vertex, [Edge])]
forall a b. (a -> b) -> a -> b
$ (Edge -> Bool) -> [Edge] -> [Edge]
forall a. (a -> Bool) -> [a] -> [a]
filter Edge -> Bool
crosses ([Edge] -> [Edge]) -> [Edge] -> [Edge]
forall a b. (a -> b) -> a -> b
$ Set Edge -> [Edge]
forall a. Set a -> [a]
Set.toList Set Edge
es
  where crosses :: Edge -> Bool
crosses (Edge (Vertex
u, Vertex
v)) = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Vertex -> Set Vertex -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Vertex
u Set Vertex
vs Bool -> Bool -> Bool
&& Vertex -> Set Vertex -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Vertex
v Set Vertex
vs
        weightEdge :: Edge -> (Vertex, [Edge])
weightEdge Edge
e = (Vertex -> Edge -> Map Edge Vertex -> Vertex
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault Vertex
0 Edge
e Map Edge Vertex
weights, [Edge
e])

-- For the two vertexes in an edge, return the vertex outside the partially constructed tree.
newVertex :: Edge -> Set Vertex -> Vertex
newVertex :: Edge -> Set Vertex -> Vertex
newVertex (Edge (Vertex
u, Vertex
v)) Set Vertex
vs
  | Vertex -> Set Vertex -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Vertex
u Set Vertex
vs = Vertex
v
  | Bool
otherwise       = Vertex
u