TP1.hs

import qualified Data.List
import qualified Data.Array
import qualified Data.Bits

-- PFL 2024/2025 Practical assignment 1

-- Uncomment the some/all of the first three lines to import the modules, do not change the code of these lines.

type City = String -- Vertex
type Path = [City]
type Distance = Int
type Edge = (City,City,Distance)

type RoadMap = [(City,City,Distance)] -- List of Edges
type AdjList = [(City,[(City,Distance)])] -- Adjacency list
type AdjMatrix = Data.Array.Array (Int,Int) (Maybe Distance) -- Matrix of distances

type Bitmask = Int

-- MinHeap data structure implemented as an array - guarantee O(1) swap operation (instead of O(n) in a list)
-- Used in Dijkstra's algorithm as a priority queue to store the vertices and their distances - represented as an array of tuples (vertex, distance)
data MinHeap a = MinHeap (Data.Array.Array Int a) Int Int -- Array, current size, capacity
-- Create an empty MinHeap
-- Time complexity: O(1)
emptyMinHeap :: (Ord a) => MinHeap a
emptyMinHeap = MinHeap (Data.Array.listArray (0, 1) []) 0 1

-- Brief: Insert an element into the heap
-- Description: Insertion's average time complexity is amortized O(1), as the array is resized only when it's full (less frequently)
-- Time complexity: O(n + complexity of heapifyUp), in the worst case
insert :: (Ord a) => a -> MinHeap a -> MinHeap a
insert x (MinHeap arr currSize capacity)
  | currSize == capacity = insert x (MinHeap newArr currSize (2*capacity)) -- O(n)
  | otherwise = MinHeap (heapifyUp currSize (arr Data.Array.// [(currSize, x)])) (currSize+1) capacity -- O(1 + complexity of heapifyUp)
  where
    -- Double the capacity of the array when it's full (to avoid resizing the array for each insertion)
    -- Fill empty space with 'undefined'
    -- Time complexity: O(n), where n is the number of elements in the heap
    newArr = Data.Array.listArray (0, 2*capacity-1) (Data.Array.elems arr ++ replicate capacity undefined)

    -- Brief: Heapify up to maintain the heap property
    -- Description:
    -- If the indexes are invalid, it returns the array as it is
    -- Inserted element at the end of the tree is compared with its parent and swapped if it's smaller
    -- In the worst case, it goes up to the root (height of the tree - log n)
    -- Time complexity: O(log n), where n is the number of elements in the heap
    heapifyUp :: (Ord a) => Int -> Data.Array.Array Int a -> Data.Array.Array Int a
    heapifyUp i arr
        | i <= 0 = arr
        | parent < 0 = arr
        | arr Data.Array.! i < arr Data.Array.! parent = heapifyUp parent (swap arr i parent)
        | otherwise = arr
      where
        parent = (i-1) `div` 2 -- parent's index

-- Brief: Extract the minimum element from the heap
-- Time complexity: O(complexity of heapifyDown)
extractMin :: (Ord a) => MinHeap a -> Maybe (a, MinHeap a)
extractMin (MinHeap arr currSize capacity)
  | currSize == 0 = Nothing -- Empty heap
  | otherwise = Just (minElem, MinHeap (heapifyDown 0 newArr) (currSize-1) capacity) -- Extract the root element and heapify down O(1 + complexity of heapifyDown)
  where
    minElem = arr Data.Array.! 0
    newArr = arr Data.Array.// [(0, arr Data.Array.! (currSize-1))] -- Replace the root with the last element O(1)

    -- Brief: Heapify down to maintain the heap property
    -- Description:
    -- The new root is swapped with the smallest child until it's smaller than both children
    -- In the worst case, it goes down to the leaf (height of the tree - log n)
    -- Time complexity: O(log n), where n is the number of elements in the heap,
    heapifyDown :: (Ord a) => Int -> Data.Array.Array Int a -> Data.Array.Array Int a
    heapifyDown i arr
      | left < currSize-1 && arr Data.Array.! left < arr Data.Array.! i = heapifyDown left (swap arr i left)
      | right < currSize-1 && arr Data.Array.! right < arr Data.Array.! i = heapifyDown right (swap arr i right)
      | otherwise = arr
      where
        left = 2*i+1 -- left child's index
        right = 2*i+2 -- right child's index

-- Brief: Swap two elements in a list
-- Time complexity: O(1)
swap :: (Ord a) => Data.Array.Array Int a -> Int -> Int -> Data.Array.Array Int a
swap arr i j = arr Data.Array.// [(i, arr Data.Array.! j), (j, arr Data.Array.! i)]



-- Brief: Returns all the cities in the graph.
-- Description:
-- 1. Extract all cities from the graph variable. Each edge has 2 cities - O(E), where E is the number of edges
-- 2. Concatenate the list of city pairs into a single list of cities - O(E)
-- 3. Sort the list of cities (with duplicates) - O(2E log(2E))
-- 4. Remove duplicates from the sorted list using the 'rmDup' function - O(2E)
-- Therefore, the sorting step is the most time-consuming step
-- Time complexity: O(E log E)
-- Note: Previous solution used Data.List.nub to remove duplicates after extracting all cities from the graph variable, which is O((2E)^2)
cities :: RoadMap -> [City]
cities graph = rmDup . Data.List.sort $ concat [[c1,c2] | (c1,c2,_) <- graph]
  where
    -- Brief: Removes duplicates from a sorted list of cities
    -- Time complexity: O(n), where n is the number of elements in the list
    rmDup :: [City] -> [City]
    rmDup [] = []
    rmDup (x:xs) = x : rmDup (dropWhile (== x) xs)

-- Brief: Returns a boolean indicating whether two cities are linked directly.
-- Description: For each edge, it checks if the given cities match both vertices.
-- Time complexity: O(E) - traversal of 'any'
areAdjacent :: RoadMap -> City -> City -> Bool
areAdjacent graph v1 v2 = any (\(c1,c2,_) -> (v1 == c1 && v2 == c2) || (v1 == c2 && v2 == c1)) graph

-- Brief: Returns a Just value with the distance between two cities connected directly, given two city names, and Nothing otherwise.
-- Description: Uses 'Data.List.find' to locate the first edge (lazy evaluation) connecting c1 and c2 - in either direction
-- Time complexity: O(E) - traversal of 'find'
distance :: RoadMap -> City -> City -> Maybe Distance
distance graph v1 v2 = case Data.List.find (\(c1,c2,_) -> (v1 == c1 && v2 == c2) || (v1 == c2 && v2 == c1)) graph of
                       Nothing -> Nothing
                       Just (_,_,dist) -> Just dist

-- Brief: Returns the cities adjacent to a given city (i.e. cities with a direct edge between them) and the respective distances to them
-- Description: For each edge in the graph variable, it checks if v matches either vertex
-- Returns an empty list if the city is not found in any edge
-- Time complexity: O(E) - traversal of the graph variable
adjacent :: RoadMap -> City -> [(City,Distance)]
adjacent graph v = [(if v == c1 then c2 else c1, d) | (c1,c2,d) <- graph, v == c1 || v == c2]

-- Brief: Returns the sum of all individual distances in a path between two cities in a Just value, if all the consecutive pairs of cities are directly connected by roads. Otherwise, it returns a Nothing.
-- Description:
-- Uses 'foldl' to accumulate the total distance by iterating over pairs of consecutive cities in the path.
-- 'accPath' checks if the distance between each pair of cities exists using 'distance'
-- Time complexity: O(V*E)
--  - As there are V-1 pairs in a path of V cities, foldl calls 'accPath' V-1 times
--  - accPath calls 'distance' for each pair of cities - O(E)
--  - At worst, the path goes through all the vertices
pathDistance :: RoadMap -> Path -> Maybe Distance
pathDistance _ [] = Just 0
pathDistance _ [_] = Just 0
pathDistance graph path = foldl accPath (Just 0) (zip path (drop 1 path))
  where
    -- Description: Accumulates the total distance between two cities in the path
    -- If any pair of cities is not directly connected, it returns Nothing
    -- Otherwise, it sums up the distances and returns the total distance (using the 'total' accumulator)
    accPath :: Maybe Distance -> (City,City) -> Maybe Distance
    accPath Nothing _ = Nothing -- if any of the edges doesn't exist (distance call returns Nothing) -> the path returns Nothing
    accPath (Just total) (v1,v2) = case distance graph v1 v2 of
                                   Nothing -> Nothing
                                   Just currDist -> Just (total + currDist)

-- Brief: Returns the names of the cities with the highest number of roads connecting to them (i.e. the vertices with the highest degree).
-- Description:
-- 1. Extracts all city occurrences from the edges in the roadmap, each edge has 2 cities
-- 2. Counts the number of occurrences for each city, these ocurrences represent the number of edges connected to each city - outdegree and indegree (undirected graph)
-- 3. Gets the maximum count among all cities
-- 4. Filters the cities with the maximum count
-- 5. Gets the city names
-- Time complexity: O(V*E) - where V is the number of cities
--  - Extracting city instances: O(E) - 'cityInstances'
--  - Getting a list of cities (no duplicates) - 'cities': O(E log E) - cityCounts
--  - Counting occurrences (traverse 'cityInstances') for each city: O(V*E) - cityCounts
--  - Getting the city counts: O(V) - maxCount
--  - Getting the maximum count: O(V) - maxCount
--  - Filtering the cities with the maximum count: O(V)
--  - Getting the city names: O(V)
rome :: RoadMap -> [City]
rome graph = map fst (filter (\(_,count) -> count == maxCount) cityCounts)
  where
    cityInstances = concat [[c1,c2] | (c1,c2,_) <- graph]
    cityCounts = [(city, length (filter (== city) cityInstances)) | city <- cities graph] -- Count the number of occurences for each city
    maxCount = maximum $ map snd cityCounts -- Get the maximum degree among all cities

-- Brief: Returns a boolean indicating whether all the cities in the graph are connected in the roadmap (i.e., if every city is reachable from every other city)
-- Description:
-- 1. For a non-empty roadmap, it performs a depth-first search (DFS) starting from an arbitrary city (the first city in the first edge, in this case)
--    Uses a bitmask to track visited cities and checks if all cities have been visited
-- 2. The 'allVisited' helper function compares the visited bitmask with the expected bitmask where all cities are marked as visited
-- Time complexity: O(V+E)
isStronglyConnected :: RoadMap -> Bool
isStronglyConnected [] = False
isStronglyConnected graph@((x,_,_):xs) = allVisited (dfs graph x (zip allCities [0..]) 0)
  where
    allCities = cities graph
    allVisited :: Bitmask -> Bool
    allVisited visited = visited == (1 `Data.Bits.shiftL` (length allCities) - 1)

-- Brief: Depth-first search algorithm to traverse the graph and mark visited cities using a bitmask
-- Time complexity: O(V+E)
dfs :: RoadMap -> City -> [(City, Int)] -> Bitmask -> Bitmask
dfs graph currC allCitiesIndex visited
  | Data.Bits.testBit visited i = visited -- City already visited
  | otherwise = foldl (\updatedVisited (adjC, _) -> dfs graph adjC allCitiesIndex updatedVisited) (Data.Bits.setBit visited i) (adjacent graph currC)
  where
    i = case Data.List.find (\(x,_) -> x == currC) allCitiesIndex of
        Just (_,index) -> index
        Nothing -> error "City not found"


-- Brief: Computes all shortest paths connecting the two cities given as input. Note that there may be more than one path with the same total distance.
--        If there are no paths between the input cities, then return an empty list. Note that the (only) shortest path between a city c and itself is [c].
-- Description: Uses Dijkstra's algorithm to find all the shortest paths between two cities
-- Detailed explanation in README.md
-- Time complexity: O(complexity of 'dijkstra')
shortestPath :: RoadMap -> City -> City -> [Path]
shortestPath graph source dest
  | source == dest = [[source]]
  | otherwise = dijkstra graph source dest

-- Brief: Dijkstra's algorithm to find all the shortest paths between two cities
-- Time complexity: O(V*E + E log V + P*V^2), where P is the number of paths
dijkstra :: RoadMap -> City -> City -> [Path]
dijkstra graph source dest = reconstructPaths endPred destIdx
  where
    allCities = cities graph -- O(E log E)
    allCitiesIndex = zip allCities [0..] -- O(V)
    adjList = [(city, adjacent graph city) | city <- allCities] -- Create the adjacency list from the list of edges - O(V*complexity of 'adjacent') == O(V*E)

    srcIdx = lookupIndex source
    destIdx = lookupIndex dest

    -- Initialize the heap, distances, and predecessors
    initPQ = insert (0, srcIdx) emptyMinHeap -- Priority queue with the source city -> O(log V)
    initDist = Data.Array.listArray (0, length allCities-1) (repeat maxBound) Data.Array.// [(srcIdx, 0)] -- Distance to each city (initialized with infinity) -> O(V)
    initPred = Data.Array.listArray (0, length allCities-1) (repeat []) -- List of predecessors of each city (initialized with empty list) -> O(V)

    (_, endPred) = dijkstraLoop adjList initPQ initDist initPred

    -- Brief: Looks up the index of a city in the list of all cities
    -- Time complexity: O(V)
    lookupIndex :: City -> Int
    lookupIndex city = case lookup city allCitiesIndex of
      Just index -> index
      Nothing -> error "City not found"

    -- Brief: Extracts the value from a Maybe type, returning a default value if it's Nothing
    -- Time complexity: O(1)
    fromMaybe :: Maybe a -> a -> a
    fromMaybe (Just x) _ = x
    fromMaybe Nothing defVal = defVal

    -- Brief: Main loop of Dijkstra's algorithm
    -- Time complexity: O(V^2 + E log V)
    dijkstraLoop :: AdjList -> MinHeap (Distance,Int) -> Data.Array.Array Int Distance -> Data.Array.Array Int [Int] -> (Data.Array.Array Int Distance, Data.Array.Array Int [Int])
    dijkstraLoop adjList pq@(MinHeap _ currSize _) distances predecessors
      | currSize == 0 = (distances, predecessors)
      | otherwise = dijkstraLoop adjList newPQ newDistances newPredecessors
      where
        Just ((dist,cIndex), xsPQ) = extractMin pq -- O(log V)
        neighbors = lookup (allCities !! cIndex) adjList -- O(V^2), as 'lookup' in adjList is O(V) and '!!' is O(V) in allCities because in Haskell lists are linked lists
        (newPQ, newDistances, newPredecessors) = foldl (\(pq, distances, predecessors) (neighbor, edgeDist) -> edgeRelax cIndex dist (pq, distances, predecessors) (neighbor, edgeDist)) (xsPQ, distances, predecessors) (fromMaybe neighbors []) -- O(complexity of 'edgeRelax' for each neighbor)

    -- Brief: Relaxes the edge between the current city and its neighbor
    -- Time complexity: O(log V), when the distance is updated
    edgeRelax :: Int -> Distance -> (MinHeap (Distance,Int), Data.Array.Array Int Distance, Data.Array.Array Int [Int]) -> (City,Distance) -> (MinHeap (Distance,Int), Data.Array.Array Int Distance, Data.Array.Array Int [Int])
    edgeRelax currIdx dist (pq, distances, predecessors) (neigh, weight)
      | newDist < distances Data.Array.! neighIdx = (insert (newDist, neighIdx) pq, distances Data.Array.// [(neighIdx, newDist)], predecessors Data.Array.// [(neighIdx, [currIdx])]) -- O(log V) for PQ insert
      | newDist == distances Data.Array.! neighIdx = (pq, distances, predecessors Data.Array.// [(neighIdx, currIdx : predecessors Data.Array.! neighIdx)]) -- O(1)
      | otherwise = (pq, distances, predecessors) -- O(1)
      where
        neighIdx = lookupIndex neigh
        newDist = dist + weight

    -- Brief: Reconstructs the paths from the destination to the source
    -- Time complexity: O(P*V^2), where P is the number of paths
    --  - Reconstructing all paths is O(P*V), where each path can have at most V cities
    --  - Converting the index paths to the city paths is O(P*V^2)
    --  - Reversing all paths is O(P*V)
    --  - P can be exponential in the worst case (complete graph)
    reconstructPaths :: Data.Array.Array Int [Int] -> Int -> [Path]
    reconstructPaths predecessors destIndex = map (reverse . indexToCity) (reconstruct destIndex)
      where
        -- Brief: Converts the index path to the city path
        -- Time complexity: O(V^2)
        --  - '!!' is O(V)
        --  - Length of the path can be at most V
        indexToCity :: [Int] -> Path
        indexToCity indexPath = map (\i -> allCities !! i) indexPath

        -- Brief: Reconstructs the paths from the destination to the source
        -- Time complexity: O(P*V), where P is the number of paths
        reconstruct :: Int -> [[Int]]
        reconstruct v
          | v == srcIdx = [[srcIdx]]
          | otherwise = [v:path | u <- predecessors Data.Array.! v, path <- reconstruct u]


-- Brief: Given a roadmap, returns a solution of the Traveling Salesperson Problem (TSP).
--        In this problem, a traveling salesperson has to visit each city exactly once and come back to the starting town. The problem is to find the shortest route, that is, the route whose total distance is minimum.
--        This problem has a known solution using dynamic programming. Any optimal TSP path will be accepted and the function only needs to return one of them, so the starting city (which is also the ending city) is left to be chosen by each group.
-- Note: The roadmap might not be a complete graph (i.e. a graph where all vertices are connected to all other vertices). If the graph does not have a TSP path, then return an empty list.
-- Description: Uses Dynamic Programming with Backtracking to find one of the optimal paths of the Traveling Salesperson Problem
-- Detailed explanation in README.md
-- Time complexity: O(2^V * V^2)
travelSales :: RoadMap -> Path
travelSales graph
  | n <= 2 = []
  | otherwise = map (cityList !!) (findPathByIndex dp)
  where
    cityList = cities graph
    n = length cityList
    adjMatrix = graphToMatrix graph
    dp = Data.Array.array ((0, 0), (2^n - 1, n-1)) [((mask, i), fillTable mask i) | mask <- [0..(2^n - 1)], i <- [0..(n-1)]] -- dp[mask][i]: Starting at city i (where we are right now), gets the minimum cost to visit all cities set in mask, ending at city with index 0

    -- Brief: Converts Roadmap to Adjacency Matrix
    -- Time complexity: O(V^2)
    graphToMatrix :: RoadMap -> AdjMatrix
    graphToMatrix graph = Data.Array.array ((0, 0), (n-1, n-1)) distances
      where distances = [((i, j), if i == j then Just 0 else distance graph (cityList !! i) (cityList !! j)) | i <- [0..n-1], j <- [0..n-1]]

    -- Brief: Uses backtracking to find the minimum cost of all of dp variable, no matter the mask or i.
    -- Note: it doesn't check if i actually has a path (0,i) yet
    fillTable :: Int -> Int -> Maybe Distance
    fillTable mask i =
      if mask == (1 `Data.Bits.shiftL` i) then adjMatrix Data.Array.! (i, 0) -- Base case: If theres only one city left to visit, and the city we are currently is that same city, it means we're done. Thus, what's left now is go to the end city (index 0) from this city.
      else foldl (
                  \md1 md2 -> case (md1, md2) of
                             (Just d1, Just d2) -> Just (min d1 d2)
                             (_, d) -> d
                 )
                 Nothing
                 [addMaybe (dp Data.Array.! (mask `Data.Bits.clearBit` i, j)) (adjMatrix Data.Array.! (i, j)) | j <- [0..(n-1)], Data.Bits.testBit mask j, j /= i]

    -- Maybe function to add two distances - used everwhere
    addMaybe :: Maybe Distance -> Maybe Distance -> Maybe Distance
    addMaybe (Just d1) (Just d2) = Just (d1 + d2)
    addMaybe _ _ = Nothing

    -- After we obtain the complete dp, we can choose the highest mask and index = 0. Then we will backtrack to find the path
    findPathByIndex :: Data.Array.Array (Int, Int) (Maybe Distance) -> [Int]
    findPathByIndex dp = getPath ((1 `Data.Bits.shiftL` n) - 1)  0            -- We start from city 0 just beacuse
      where                          -- ^^^^^^^^ Mask with all the cities to visit
        -- Brief: Find the path backwards
        getPath :: Int -> Int -> [Int]
        getPath mask indexNow
          | mask == (1 `Data.Bits.shiftL` indexNow)  = indexNow : [0]   -- Base case: If we have visited all the cities except city 0, then the first city to visit has indexNow adn after we will visit for sure city 0
          | otherwise =
              if null allCitiesToVisit then []
              else indexNow : getPath newMask (snd $ Data.List.minimumBy (\(dist1, _) (dist2, _) -> compare dist1 dist2) allCitiesToVisit) -- <-city to visit next
              where
                newMask = mask `Data.Bits.clearBit` indexNow
                allCitiesToVisit = [(dist, j) | j <- [0..(n-1)], Data.Bits.testBit mask j, j /= indexNow, let dist = addMaybe (adjMatrix Data.Array.! (indexNow, j)) (dp Data.Array.! (newMask, j)), dist /= Nothing]


tspBruteForce :: RoadMap -> Path
tspBruteForce = undefined -- only for groups of 3 people; groups of 2 people: do not edit this function

-- Some graphs to test your work
gTest1 :: RoadMap
gTest1 = [("7","6",1),("8","2",2),("6","5",2),("0","1",4),("2","5",4),("8","6",6),("2","3",7),("7","8",7),("0","7",8),("1","2",8),("3","4",9),("5","4",10),("1","7",11),("3","5",14)]

gTest2 :: RoadMap
gTest2 = [("0","1",10),("0","2",15),("0","3",20),("1","2",35),("1","3",25),("2","3",30)]

gTest3 :: RoadMap -- unconnected graph
gTest3 = [("0","1",4),("2","3",2)]