Prim algorithm for MST

import Data.List

-- В graph на каждом шагу будем хранить ребра, у которых хотя бы одна вершина НЕ содержиться в текущем MST (Minimal Spanning Tree)
-- Также будем поддерживать список ребёр в MST и список неприсоединённых к MST вершин
-- На каждом шагу выбираем ребро минимального веса, среди смежных с текущим MST, за начальное ребро берём наименьшее среди всех
prim :: (Ord c, Eq a) => [(a, a, c)] -> [(a, a, c)]
prim [] = []
prim graph = primH (tail sortedEdges) mst unvisited where -- Основная функция, подготавливает данные для вспомогательной
  sortedEdges = nub $ sortBy (\(_, _, w1) (_, _, w2) -> compare w1 w2) graph -- Сортируем ребра по весу и удаляем дубликаты
  mst = [head sortedEdges] -- Добавляем первое ребро в MST
  unvisited = foldl (flip delete) (nub $ concatMap getVertices graph) (getVertices $ head mst) -- Создаем список неприсоединённых к MST вершин
  primH graph mst unvisited -- Вспомогательная функция, рекурсивно выполняет шаги алгоритма
      | null graph = mst
      | null unvisited = mst
      | otherwise = primH newGraph newMst newUnvisited where
        (from, to, weight) = head (filter (\(x, y, _) -> notElem x unvisited || notElem y unvisited) graph) -- Берём ребро с наименьшим весом, соединённое с MST
        newMst = (from, to, weight) : mst -- Добавляем это ребро в MST
        newUnvisited = foldl (flip delete) unvisited [from, to] -- Удаляем вершину из непосещённых
        newGraph = filter (\(x, y, _) -> elem x newUnvisited || elem y newUnvisited) graph -- Убираем из списка ребёр, те что соединяют вершины из MST
  getVertices (x, y, _) = [x, y] -- Функция для получения списка вершин некоторого ребра

graph :: [(Int, Int, Int)]
graph = [(1, 2, 2), (1, 3, 5), (2, 3, 1), (2, 4, 6), (2, 5, 3), (3, 5, 1), (4, 5, 3)]

main :: IO ()
main = print $ prim graph