Skip to content

Data structures

AVL-Tree Map

module Map {

  import Number(max)

  type alias MapFields<key, val> =
    { key    : key
    , value  : val
    , height : int32
    , left   : Map<key, val>
    , right  : Map<key, val> 
    }

  type Map<key, val> 
    = Empty 
    | Map(MapFields<key, val>)

  // Helpers

  fun height(map) =
    match(map) {
      | Empty => 0
      | Map({ height = h | _ }) => h
    }

  fun make_node(key, val, left, right) =
    Map({ key    = key
        , value  = val
        , height = 1 + max(height(left), height(right))
        , left   = left
        , right  = right
        })

  fun balance_factor
    | Empty  => 0
    | Map(m) => height(m.left) - height(m.right)

  // Constructors

  let empty = Empty

  fun singleton(key, val) = 
    make_node(key, val, Empty, Empty)

  // Rotations

  //
  //       y              x
  //      / \            / \
  //     x   t3   ==>   t1  y
  //    / \                / \
  //   t1 t2              t2 t3
  //

  fun rotate_right(map) = 
    match(map) {
      | Map({ key = yk, value = yv, left = Map(x), right = t3 | _ }) => 
          let t1 = x.left;
              t2 = x.right
          in
          make_node(x.key, x.value, t1, make_node(yk, yv, t2, t3))
      | _ => 
          map
    }

  //
  //     x                 y
  //    / \               / \
  //   t1  y     ==>     x  t3
  //      / \           / \
  //     t2 t3         t1 t2
  //

  fun rotate_left(map) = 
    match(map) {
      | Map({ key = xk, value = xv, left = t1, right = Map(y) | _ }) => 
          let t2 = y.left;
              t3 = y.right
          in
          make_node(y.key, y.value, make_node(xk, xv, t1, t2), t3)
      | _ => 
          map
    }

  // Rebalancing

  fun rebalance(map) =
    match(map) {
      | Empty =>
          Empty

      | Map(m) =>
          let bf = balance_factor(map) in
          if (bf > 1) then
            if (balance_factor(m.left) < 0) then
              rotate_right(
                make_node(m.key, m.value,
                  rotate_left(m.left),
                  m.right
                )
              )
            else
              rotate_right(map)

          else if (bf < -1) then
            if (balance_factor(m.right) > 0) then
              rotate_left(
                make_node(m.key, m.value,
                  m.left,
                  rotate_right(m.right)
                )
              )
            else
              rotate_left(map)

          else
            map
    }

  // Lookup

  fun lookup(key, map) =
    fold(map) {
      | Empty =>
          None

      | Map({ key = key_, value = value, left = @left, right = @right | _ }) =>
          if (key == key_) then
            Some(value)
          else 
            if (key < key_) then
              left
            else
              right
    }

  // Insertion

  fun insert(key, val, map) =
    fold(map) {
      | Empty =>
          singleton(key, val)

      | Map({ left = @left, right = @right | _ } as m) =>
          if (key == m.key) then
            make_node(key, val, m.left, m.right)

          else if (key < m.key) then
            rebalance(make_node(m.key, m.value, left, m.right))

          else
            rebalance(make_node(m.key, m.value, m.left, right))
    }

  // Deletion

  fold delete_min_node : Map<k, v> -> (Option<k>, Option<v>, Map<k, v>) {
    | Empty =>
        (None, None, Empty)

    | Map({ left = Empty | _ } as m) =>
        (Some(m.key), Some(m.value), m.right)

    | Map({ left = delete_min_node(@left) | _ } as m) =>
        let (k, v, new_left) = left 
        in
        (k, v, rebalance(make_node(m.key, m.value, new_left, m.right)))
  }

  fun delete(key, map) =
    fold(map) {
      | Empty =>
          Empty

      | Map({ left = @left, right = @right | _ } as m) =>
          if (key < m.key) then
            rebalance(make_node(m.key, m.value, left, m.right))

          else if (key > m.key) then
            rebalance(make_node(m.key, m.value, m.left, right))

          else
            match(delete_min_node(m.right)) {
              | (Some(k), Some(v), new_right) =>
                  rebalance(make_node(k, v, m.left, new_right))
              | _ =>
                  m.left
            }
    }
}module Map {

  import Number(max)

  type alias MapFields<key, val> =
    { key    : key
    , value  : val
    , height : int32
    , left   : Map<key, val>
    , right  : Map<key, val> 
    }

  type Map<key, val> 
    = Empty 
    | Map(MapFields<key, val>)

  // Helpers

  fun height(map) =
    match(map) {
      | Empty => 0
      | Map({ height = h | _ }) => h
    }

  fun make_node(key, val, left, right) =
    Map({ key    = key
        , value  = val
        , height = 1 + max(height(left), height(right))
        , left   = left
        , right  = right
        })

  fun balance_factor
    | Empty  => 0
    | Map(m) => height(m.left) - height(m.right)

  // Constructors

  let empty = Empty

  fun singleton(key, val) = 
    make_node(key, val, Empty, Empty)

  // Rotations

  //
  //       y              x
  //      / \            / \
  //     x   t3   ==>   t1  y
  //    / \                / \
  //   t1 t2              t2 t3
  //

  fun rotate_right(map) = 
    match(map) {
      | Map({ key = yk, value = yv, left = Map(x), right = t3 | _ }) => 
          let t1 = x.left;
              t2 = x.right
          in
          make_node(x.key, x.value, t1, make_node(yk, yv, t2, t3))
      | _ => 
          map
    }

  //
  //     x                 y
  //    / \               / \
  //   t1  y     ==>     x  t3
  //      / \           / \
  //     t2 t3         t1 t2
  //

  fun rotate_left(map) = 
    match(map) {
      | Map({ key = xk, value = xv, left = t1, right = Map(y) | _ }) => 
          let t2 = y.left;
              t3 = y.right
          in
          make_node(y.key, y.value, make_node(xk, xv, t1, t2), t3)
      | _ => 
          map
    }

  // Rebalancing

  fun rebalance(map) =
    match(map) {
      | Empty =>
          Empty

      | Map(m) =>
          let bf = balance_factor(map) in
          if (bf > 1) then
            if (balance_factor(m.left) < 0) then
              rotate_right(
                make_node(m.key, m.value,
                  rotate_left(m.left),
                  m.right
                )
              )
            else
              rotate_right(map)

          else if (bf < -1) then
            if (balance_factor(m.right) > 0) then
              rotate_left(
                make_node(m.key, m.value,
                  m.left,
                  rotate_right(m.right)
                )
              )
            else
              rotate_left(map)

          else
            map
    }

  // Lookup

  fun lookup(key, map) =
    fold(map) {
      | Empty =>
          None

      | Map({ key = key_, value = value, left = @left, right = @right | _ }) =>
          if (key == key_) then
            Some(value)
          else 
            if (key < key_) then
              left
            else
              right
    }

  // Insertion

  fun insert(key, val, map) =
    fold(map) {
      | Empty =>
          singleton(key, val)

      | Map({ left = @left, right = @right | _ } as m) =>
          if (key == m.key) then
            make_node(key, val, m.left, m.right)

          else if (key < m.key) then
            rebalance(make_node(m.key, m.value, left, m.right))

          else
            rebalance(make_node(m.key, m.value, m.left, right))
    }

  // Deletion

  fold delete_min_node : Map<k, v> -> (Option<k>, Option<v>, Map<k, v>) {
    | Empty =>
        (None, None, Empty)

    | Map({ left = Empty | _ } as m) =>
        (Some(m.key), Some(m.value), m.right)

    | Map({ left = delete_min_node(@left) | _ } as m) =>
        let (k, v, new_left) = left 
        in
        (k, v, rebalance(make_node(m.key, m.value, new_left, m.right)))
  }

  fun delete(key, map) =
    fold(map) {
      | Empty =>
          Empty

      | Map({ left = @left, right = @right | _ } as m) =>
          if (key < m.key) then
            rebalance(make_node(m.key, m.value, left, m.right))

          else if (key > m.key) then
            rebalance(make_node(m.key, m.value, m.left, right))

          else
            match(delete_min_node(m.right)) {
              | (Some(k), Some(v), new_right) =>
                  rebalance(make_node(k, v, m.left, new_right))
              | _ =>
                  m.left
            }
    }
}