(* Coder: Rafał Łyżwa
 * Reviewer: Marek Żochowski *)

type ('k, 'v) map =
  | Empty
  | Node of ('k, 'v) map * 'k * 'v * ('k, 'v) map * int

type ('k, 'v) t =
  {
    cmp : 'k -> 'k -> int;
    map : ('k, 'v) map;
  }

let height = function
  | Node (_, _, _, _, h) -> h
  | Empty -> 0

let make l k v r = Node (l, k, v, r, max (height l) (height r) + 1)

let bal l k v r =
  let hl = height l in
  let hr = height r in
  if hl > hr + 2 then
    match l with
    | Node (ll, lk, lv, lr, _) ->
        if height ll >= height lr then make ll lk lv (make lr k v r)
        else
          (match lr with
          | Node (lrl, lrk, lrv, lrr, _) ->
              make (make ll lk lv lrl) lrk lrv (make lrr k v r)
          | Empty -> assert false)
    | Empty -> assert false
  else if hr > hl + 2 then
    match r with
    | Node (rl, rk, rv, rr, _) ->
        if height rr >= height rl then make (make l k v rl) rk rv rr
        else
          (match rl with
          | Node (rll, rlk, rlv, rlr, _) ->
              make (make l k v rll) rlk rlv (make rlr rk rv rr)
          | Empty -> assert false)
    | Empty -> assert false
  else Node (l, k, v, r, max hl hr + 1)

let rec min_binding = function
  | Node (Empty, k, v, _, _) -> k, v
  | Node (l, _, _, _, _) -> min_binding l
  | Empty -> raise Not_found

let rec remove_min_binding = function
  | Node (Empty, _, _, r, _) -> r
  | Node (l, k, v, r, _) -> bal (remove_min_binding l) k v r
  | Empty -> invalid_arg "PMap.remove_min_binding"

let merge t1 t2 =
  match t1, t2 with
  | Empty, _ -> t2
  | _, Empty -> t1
  | _ ->
      let k, v = min_binding t2 in
      bal t1 k v (remove_min_binding t2)

let create cmp = { cmp = cmp; map = Empty }
let empty = { cmp = compare; map = Empty }

let is_empty x =
    x.map = Empty

let add x d { cmp = cmp; map = map } =
  let rec loop = function
    | Node (l, k, v, r, h) ->
        let c = cmp x k in
        if c = 0 then Node (l, x, d, r, h)
        else if c < 0 then
          let nl = loop l in
          bal nl k v r
        else
          let nr = loop r in
          bal l k v nr
    | Empty -> Node (Empty, x, d, Empty, 1) in
  { cmp = cmp; map = loop map }

let find x { cmp = cmp; map = map } =
  let rec loop = function
    | Node (l, k, v, r, _) ->
        let c = cmp x k in
        if c < 0 then loop l
        else if c > 0 then loop r
        else v
    | Empty -> raise Not_found in
  loop map

let remove x { cmp = cmp; map = map } =
  let rec loop = function
    | Node (l, k, v, r, _) ->
        let c = cmp x k in
        if c = 0 then merge l r else
        if c < 0 then bal (loop l) k v r else bal l k v (loop r)
    | Empty -> Empty in
  { cmp = cmp; map = loop map }

let mem x { cmp = cmp; map = map } =
  let rec loop = function
    | Node (l, k, v, r, _) ->
        let c = cmp x k in
        c = 0 || loop (if c < 0 then l else r)
    | Empty -> false in
  loop map

let exists = mem

let iter f { map = map } =
  let rec loop = function
    | Empty -> ()
    | Node (l, k, v, r, _) -> loop l; f k v; loop r in
  loop map

let map f { cmp = cmp; map = map } =
  let rec loop = function
    | Empty -> Empty
    | Node (l, k, v, r, h) ->
      let l = loop l in
      let r = loop r in
      Node (l, k, f v, r, h) in
  { cmp = cmp; map = loop map }

let mapi f { cmp = cmp; map = map } =
  let rec loop = function
    | Empty -> Empty
    | Node (l, k, v, r, h) ->
      let l = loop l in
      let r = loop r in
      Node (l, k, f k v, r, h) in
  { cmp = cmp; map = loop map }

let fold f { cmp = cmp; map = map } acc =
  let rec loop acc = function
    | Empty -> acc
    | Node (l, k, v, r, _) ->
      loop (f v (loop acc l)) r in
  loop acc map

let foldi f { cmp = cmp; map = map } acc =
  let rec loop acc = function
    | Empty -> acc
    | Node (l, k, v, r, _) ->
       loop (f k v (loop acc l)) r in
  loop acc map

exception Cykliczne

let szukaj graf =
(* mapa trzymająca stopnie wchodzące wierzchołków    *)
  let stopien = ref empty in
(* mapa trzymająca listy krawędzi wychodzących       *)
  let krawedzie = ref empty in
(* lista wierzchołków o zerowym stopniu wchodzącym   *)
  let zerowe = ref [] in
(* lista wierzchołków posortowanych topologicznie    *)
  let posortowane = ref [] in
(* lista zwracana jako wynik                         *)
  let wynik = ref [] in
(* ilość wierzchołków w grafie                       *)
  let ile = ref 0 in
  let sprawdzczyzerowy wierzcholek =
    if find wierzcholek !stopien = 0 then
      zerowe := wierzcholek :: !zerowe in
  let powieksz wierzcholek =
    stopien := (add wierzcholek (find wierzcholek !stopien + 1) !stopien) in
  let zmniejsz wierzcholek =
    stopien := (add wierzcholek (find wierzcholek !stopien - 1) !stopien);
    sprawdzczyzerowy wierzcholek in
(* utwórz w mapach pola dotyczące danego wierzchołka *)
  let dotknij wierzcholek =
    if mem wierzcholek !krawedzie = false then
      krawedzie := (add wierzcholek [] !krawedzie);
    if mem wierzcholek !stopien = false then
      stopien := (add wierzcholek 0 !stopien) in
  let przetworz (aktualny, listakrawedzi) =
    dotknij aktualny;
    krawedzie := (add aktualny (listakrawedzi@(find aktualny !krawedzie)) !krawedzie);
    List.iter dotknij listakrawedzi;
    List.iter powieksz listakrawedzi in
  let usun wierzcholek =
    posortowane := wierzcholek :: !posortowane;
    List.iter zmniejsz (find wierzcholek !krawedzie) in
  let aktualny = ref (fst (List.hd graf)) in
  List.iter przetworz graf;
  iter (fun x y -> sprawdzczyzerowy x) !stopien;
  while (!zerowe <> [])
  do
    aktualny := (List.hd !zerowe);
    zerowe := List.tl !zerowe;
    usun !aktualny
  done;
  iter (fun x y -> ile := !ile + 1) !stopien;
  if List.length !posortowane < !ile then
    raise Cykliczne
  else
    wynik := (List.rev !posortowane);
  !wynik

let topol graf =
  match graf with
  | [] -> []
  | h::t -> szukaj graf