priorityqueue.gr

/**
 * A priority queue is a data structure that maintains elements in a priority order. Elements with higher priority are served before elements with lower priority when extracting from the priority queue.
 *
 * An immutable priority queue implementation is available in the `Immutable` submodule.
 *
 * @example from "priorityqueue" include Priorityqueue
 *
 * @since v0.5.3
 */
module PriorityQueue

from "array" include Array
from "list" include List
from "number" include Number
from "option" include Option

/**
 * Mutable data structure which maintains a priority order for its elements.
 */
abstract record PriorityQueue<a> {
  mut size: Number,
  mut array: Array<Option<a>>,
  comp: (a, a) => Number,
}

let swap = (i1, i2, array) => {
  let t = array[i2]
  array[i2] = array[i1]
  array[i1] = t
}

let get = (array, i) =>
  Option.expect(
    "Impossible: "
      ++ toString(i)
      ++ " in PriorityQueue's inner storage array is None",
    array[i]
  )

let rec siftDown = (i, pq) => {
  let leftI = 2 * i + 1
  let rightI = 2 * i + 2

  // we want to find the smaller child from the current tree node to sift down to
  let mut swapWithI = i
  if (leftI < pq.size && pq.comp(get(pq.array, leftI), get(pq.array, i)) < 0) {
    swapWithI = leftI
  }
  if (
    rightI < pq.size
    && pq.comp(get(pq.array, rightI), get(pq.array, swapWithI)) < 0
  ) {
    swapWithI = rightI
  }
  if (swapWithI != i) {
    swap(i, swapWithI, pq.array)
    siftDown(swapWithI, pq)
  }
}

let rec siftUp = (i, pq) => {
  let parentI = Number.trunc((i - 1) / 2)
  // we should only sift up if the element is smaller than its parent
  if (i > 0 && pq.comp(get(pq.array, i), get(pq.array, parentI)) < 0) {
    swap(i, parentI, pq.array)
    siftUp(parentI, pq)
  }
}

/**
 * Creates a new priority queue with a given internal storage size and a
 * comparator function, which is used to determine priority of elements. The
 * comparator function takes two elements and must return 0 if both share
 * priority, a positive number if the first has greater priority, and a
 * negative number if the first has less priority.
 *
 * Generally, you won't need to care about the storage size of your priority
 * queue and can use the default size.
 *
 * @param compare: The comparator function used to indicate priority order
 * @param size: The initial storage size of the priority queue
 * @returns An empty priority queue
 *
 * @since v0.5.3
 * @history v0.6.0: Merged with `makeSized`; modified signature to accept size
 *
 * @example PriorityQueue.make() // creates a min priority queue of numbers using the compare pervasive
 * @example PriorityQueue.make(compare=compare, size=32) // creates a min priority queue of numbers using the compare pervasive and an initial size of 32
 * @example PriorityQueue.make((a, b) => String.length(b) - String.length(a)) // creates a priority queue by string length (longest to shortest)
 */
provide let make = (compare=compare, size=16) => {
  { size: 0, array: Array.make(size, None), comp: compare }
}

/**
 * Gets the number of elements in a priority queue.
 *
 * @param pq: The priority queue to inspect
 * @returns The number of elements in the priority queue
 *
 * @since v0.5.3
 */
provide let size = pq => {
  pq.size
}

/**
 * Determines if the priority queue contains no elements.
 *
 * @param pq: The priority queue to check
 * @returns `true` if the priority queue is empty and `false` otherwise
 *
 * @since v0.5.3
 */
provide let isEmpty = pq => {
  pq.size == 0
}

/**
 * Adds a new element to the priority queue.
 *
 * @param val: The value to add into the priority queue
 * @param pq: The priority queue to update
 *
 * @since v0.5.3
 */
provide let push = (val, pq) => {
  let arrLen = Array.length(pq.array)
  // double size of internal array if out of space
  if (pq.size == arrLen) {
    let oldArr = pq.array
    pq.array = Array.make(arrLen * 2, None)
    Array.forEachi((val, i) => {
      pq.array[i] = val
    }, oldArr)
  }
  pq.array[pq.size] = Some(val)
  pq.size += 1
  // reorder heap to ensure that binary heap property of parent nodes having
  // larger values than their children is upheld
  siftUp(pq.size - 1, pq)
}

/**
 * Retrieves the highest priority element in the priority queue. It is not
 * removed from the queue.
 *
 * @param pq: The priority queue to inspect
 * @returns `Some(value)` containing the highest priority element or `None` if the priority queue is empty
 *
 * @since v0.5.3
 */
provide let peek = pq => {
  if (pq.size == 0) {
    None
  } else {
    pq.array[0]
  }
}

/**
 * Removes and retrieves the highest priority element in the priority queue.
 *
 * @param pq: The priority queue to inspect
 * @returns `Some(value)` containing the highest priority element or `None` if the priority queue is empty
 *
 * @since v0.5.3
 */
provide let pop = pq => {
  if (pq.size == 0) return None
  let root = pq.array[0]

  pq.array[0] = pq.array[pq.size - 1]
  pq.array[pq.size - 1] = None
  pq.size -= 1
  // reorder heap to ensure that binary heap property of parent nodes having
  // larger values than their children is upheld
  siftDown(0, pq)
  return root
}

/**
 * Clears the priority queue and produces a list of all of the elements in the priority
 * queue in priority order.
 *
 * @param pq: The priority queue to drain
 * @returns A list of all elements in the priority in priority order
 *
 * @since v0.5.3
 */
provide let drain = pq => {
  let rec drainRec = acc => {
    match (pop(pq)) {
      Some(val) => drainRec([val, ...acc]),
      None => acc,
    }
  }
  List.reverse(drainRec([]))
}

/**
 * Constructs a new priority queue initialized with the elements in the array
 * using a custom comparator function, which is used to determine priority of
 * elements. The comparator function takes two elements and must return 0 if
 * both share priority, a positive number if the first has greater priority,
 * and a negative number if the first has less priority.
 *
 * @param array: An array of values used to initialize the priority queue
 * @param compare: A comparator function used to assign priority to elements
 * @returns A priority queue containing the elements from the array
 *
 * @since v0.5.4
 * @history v0.6.0: Made `compare` a default argument
 */
provide let fromArray = (array, compare=compare) => {
  let size = Array.length(array)
  let array = Array.map(x => Some(x), array)
  let heap = { size, array, comp: compare }
  for (let mut i = size - 1; i >= 0; i -= 1) {
    siftDown(i, heap)
  }
  heap
}

/**
 * Constructs a new priority queue initialized with the elements in the list
 * using a custom comparator function, which is used to determine priority of
 * elements. The comparator function takes two elements and must return 0 if
 * both share priority, a positive number if the first has greater priority,
 * and a negative number if the first has less priority.
 *
 * @param list: A list of values used to initialize the priority queue
 * @param compare: A comparator function used to assign priority to elements
 * @returns A priority queue containing the elements from the list
 *
 * @since v0.5.3
 * @history v0.6.0: Made `compare` a default argument
 */
provide let fromList = (list, compare=compare) => {
  let array = Array.fromList(list)
  fromArray(array, compare=compare)
}

/**
 * An immutable priority queue implementation.
 *
 * @since v0.6.0
 * @history v0.5.4: Originally in `"immutablepriorityqueue"` module
 */
provide module Immutable {
  // implementation based on immutable skew binomial queue with global root optimization
  // as described in the paper "Optimal Purely Functional Priority Queues" by Chris Okasaki.

  // rank is a stand-in for height of this skew binomial tree
  record rec Node<a> {
    val: a,
    rank: Number,
    children: List<Node<a>>,
  }

  // an optimization over binomial queue that allows keeping track of the root value

  // a skew binomial queue is defined as a forest of heap-ordered skew binomial trees
  record PQRoot<a> {
    rootVal: a,
    pq: List<Node<a>>,
  }

  /**
   * Immutable data structure which maintains a priority order for its elements.
   */
  abstract record PriorityQueue<a> {
    comp: (a, a) => Number,
    size: Number,
    root: Option<PQRoot<a>>,
  }

  /**
   * An empty priority queue with the default `compare` comparator.
   *
   * @since v0.6.0
   * @history v0.5.4: Originally in `"immutablepriorityqueue"` module
   */
  provide let empty = {
    let empty = { comp: compare, size: 0, root: None }
    empty
  }

  /**
   * Creates a new priority queue with a comparator function, which is used to
   * determine priority of elements. The comparator function takes two elements
   * and must return 0 if both share priority, a positive number if the first
   * has greater priority, and a negative number if the first has less priority.
   *
   * @param compare: The comparator function used to indicate priority order
   * @returns An empty priority queue
   *
   * @example PriorityQueue.Immutable.make(compare) // creates a min priority queue of numbers using the compare pervasive
   * @example PriorityQueue.Immutable.make((a, b) => String.length(b) - String.length(a)) // creates a priority queue by string length (longest to shortest)
   *
   * @since v0.6.0
   * @history v0.5.3: Originally in `"immutablepriorityqueue"` module with `compare` being a required argument
   */
  provide let make = (compare=compare) => {
    { comp: compare, size: 0, root: None }
  }

  /**
   * Gets the number of elements in a priority queue.
   *
   * @param pq: The priority queue to inspect
   * @returns The number of elements in the priority queue
   *
   * @since v0.6.0
   * @history v0.5.3: Originally in `"immutablepriorityqueue"` module
   */
  provide let size = ({ size, _ } as pq: PriorityQueue<a>) => {
    size
  }

  /**
   * Determines if the priority queue contains no elements.
   *
   * @param pq: The priority queue to check
   * @returns `true` if the priority queue is empty and `false` otherwise
   *
   * @since v0.6.0
   * @history v0.5.3: Originally in `"immutablepriorityqueue"` module
   */
  provide let isEmpty = ({ size, _ } as pq: PriorityQueue<a>) => {
    size == 0
  }

  let skewLinkNodes = (comp, newNode, node1, node2) => {
    // make the two nodes with larger values children of the node with the smallest value
    if (comp(node1.val, newNode.val) <= 0 && comp(node1.val, node2.val) <= 0) {
      {
        val: node1.val,
        rank: node1.rank + 1,
        children: [newNode, node2, ...node1.children],
      }
    } else if (
      comp(node2.val, newNode.val) <= 0
      && comp(node2.val, node1.val) <= 0
    ) {
      {
        val: node2.val,
        rank: node2.rank + 1,
        children: [newNode, node1, ...node2.children],
      }
    } else {
      { val: newNode.val, rank: node1.rank + 1, children: [node1, node2] }
    }
  }

  let skewInsert = (comp, val, pq) => {
    let newNode = { val, rank: 0, children: [] }
    match (pq) {
      // the only time two trees will have the same rank is if they are the
      // smallest-ranked trees in the queue, in which case we should link
      // them with the new node
      [node1, node2, ...rest] when node1.rank == node2.rank =>
        [skewLinkNodes(comp, newNode, node1, node2), ...rest],
      _ => [newNode, ...pq],
    }
  }

  /**
   * Produces a new priority queue by inserting the given element into the given priority queue.
   *
   * @param val: The value to add into the priority queue
   * @param pq: The priority queue
   * @returns A new priority queue with the given element inserted
   *
   * @since v0.6.0
   * @history v0.5.3: Originally in `"immutablepriorityqueue"` module
   */
  provide let push = (val, pq) => {
    let { comp, size, root } = pq
    match (root) {
      None => { comp, size: 1, root: Some({ rootVal: val, pq: [] }) },
      Some({ rootVal, pq }) => {
        // make the new value the root if it has higher priority than the highest priority value
        let (morePriorityVal, lessPriorityVal) = if (comp(val, rootVal) <= 0)
          (val, rootVal)
        else
          (rootVal, val)

        let newRoot = Some(
          {
            rootVal: morePriorityVal,
            pq: skewInsert(comp, lessPriorityVal, pq),
          },
        )
        { comp, size: size + 1, root: newRoot }
      },
    }
  }

  /**
   * Retrieves the highest priority element in the priority queue. It is not
   * removed from the queue.
   *
   * @param pq: The priority queue to inspect
   * @returns `Some(value)` containing the highest priority element or `None` if the priority queue is empty
   *
   * @since v0.6.0
   * @history v0.5.3: Originally in `"immutablepriorityqueue"` module
   */
  provide let peek = pq => {
    match (pq.root) {
      None => None,
      Some({ rootVal, _ }) => Some(rootVal),
    }
  }

  let linkNodes = (comp, node1, node2) => {
    // make the node with higher priority the parent of the node with smaller
    // priority to presere heap-ordering
    let (morePriority, lessPriority) = if (comp(node1.val, node2.val) <= 0)
      (node1, node2)
    else
      (node2, node1)
    {
      val: morePriority.val,
      rank: morePriority.rank + 1,
      children: [lessPriority, ...morePriority.children],
    }
  }

  // step through the trees in the priority queue in increasing rank order until we
  // find a missing rank, linking trees of equal rank as we go
  let rec ins = (comp, node, pq) => {
    match (pq) {
      [] => [node],
      [firstNode, ...rest] => {
        if (node.rank < firstNode.rank) {
          [node, firstNode, ...rest]
        } else {
          ins(comp, linkNodes(comp, node, firstNode), rest)
        }
      },
    }
  }

  let uniquify = (comp, pq) => {
    match (pq) {
      [] => [],
      [node, ...rest] => ins(comp, node, rest),
    }
  }

  // step through the trees of two priority queues in increasing rank order,
  // performing a simple link whenever we find two trees of equal rank
  let rec mergeUniq = (comp, pq1, pq2) => {
    match ((pq1, pq2)) {
      ([], ts) | (ts, []) => ts,
      ([node1, ...rest1], [node2, ...rest2]) => {
        if (node1.rank < node2.rank) {
          [node1, ...mergeUniq(comp, rest1, pq2)]
        } else if (node2.rank < node1.rank) {
          [node2, ...mergeUniq(comp, pq1, rest2)]
        } else {
          ins(
            comp,
            linkNodes(comp, node1, node2),
            mergeUniq(comp, rest1, rest2)
          )
        }
      },
    }
  }

  let merge = (comp, pq1, pq2) =>
    mergeUniq(comp, uniquify(comp, pq1), uniquify(comp, pq2))

  // splits the node with the minimum value out from the rest of the nodes
  let rec separateHighestPriority = (comp, pq) => {
    match (pq) {
      // empty list case should never happen; this is here just to satisfy the compiler
      [] => fail "Impossible: getHighestPriority called with empty PQ",
      [node] => (node, []),
      [node, ...rest] => {
        let (currMinNode, currNonMinNodes) = separateHighestPriority(comp, rest)
        if (comp(node.val, currMinNode.val) <= 0) {
          (node, rest)
        } else {
          (currMinNode, [node, ...currNonMinNodes])
        }
      },
    }
  }

  // splits the nodes of rank 0 out from the other nodes
  let rec splitRankZero = (rankZeroVals, nonRankZeroNodes, pq) => {
    match (pq) {
      [] => (rankZeroVals, nonRankZeroNodes),
      [node, ...rest] => {
        if (node.rank == 0) {
          splitRankZero([node.val, ...rankZeroVals], nonRankZeroNodes, rest)
        } else {
          splitRankZero(rankZeroVals, [node, ...nonRankZeroNodes], rest)
        }
      },
    }
  }

  let withoutHighestPriority = (comp, pq) => {
    // split out the node with the highest priority
    let (hpNode, nonHpNodes) = separateHighestPriority(comp, pq)
    // split out the values with nodes of rank 0, which will all be singleton nodes
    let (rankZeroVals, nonRankZeroNodes) = splitRankZero(
      [],
      [],
      hpNode.children
    )

    let mergedPq = merge(comp, nonHpNodes, nonRankZeroNodes)
    List.reduce((pq, val) => skewInsert(comp, val, pq), mergedPq, rankZeroVals)
  }

  let rec findHighestPriority = (comp, pq) => {
    match (pq) {
      // empty list case should never happen; this is here just to satisfy the compiler
      [] => fail "Impossible: findHighestPriority with empty PQ",
      [node] => node.val,
      [node, ...rest] => {
        let currMin = findHighestPriority(comp, rest)
        if (comp(node.val, currMin) <= 0) node.val else currMin
      },
    }
  }

  /**
   * Produces a new priority queue without the highest priority element in the
   * given priority queue. If the input priority queue is empty, this function will
   * return it.
   *
   * @param pq: The priority queue
   * @returns A new priority queue without the highest priority element
   *
   * @since v0.6.0
   * @history v0.5.3: Originally in `"immutablepriorityqueue"` module
   */
  provide let pop = pq => {
    let pqWithRoot = pq
    let { comp, size, root } = pq
    match (root) {
      None => pq,
      Some({ pq, rootVal }) => {
        let newRoot = if (pq == []) {
          None
        } else {
          Some(
            {
              rootVal: findHighestPriority(comp, pq),
              pq: withoutHighestPriority(comp, pq),
            },
          )
        }
        { comp, size: size - 1, root: newRoot }
      },
    }
  }

  /**
   * Produces a list of all elements in the priority queue in priority order.
   *
   * @param pq: The priority queue to drain
   * @returns A list of all elements in the priority in priority order
   *
   * @since v0.6.0
   * @history v0.5.3: Originally in `"immutablepriorityqueue"` module
   */
  provide let drain = pq => {
    let rec drainRec = (acc, pq) => {
      match (pq.root) {
        None => acc,
        Some(root) => drainRec([root.rootVal, ...acc], pop(pq)),
      }
    }
    List.reverse(drainRec([], pq))
  }

  /**
   * Constructs a new priority queue initialized with the elements in the list
   * using a custom comparator function, which is used to determine priority of
   * elements. The comparator function takes two elements and must return 0 if
   * both share priority, a positive number if the first has greater priority,
   * and a negative number if the first has less priority.
   *
   * @param list: A list of values used to initialize the priority queue
   * @param compare: A comparator function used to assign priority to elements
   * @returns A priority queue containing the elements from the list
   *
   * @since v0.6.0
   * @history v0.5.3: Originally in `"immutablepriorityqueue"` module with `compare` being a required argument
   */
  provide let fromList = (list, compare=compare) => {
    List.reduce((pq, val) => push(val, pq), make(compare=compare), list)
  }

  /**
   * Constructs a new priority queue initialized with the elements in the array
   * using a custom comparator function, which is used to determine priority of
   * elements. The comparator function takes two elements and must return 0 if
   * both share priority, a positive number if the first has greater priority,
   * and a negative number if the first has less priority.
   *
   * @param array: An array of values used to initialize the priority queue
   * @param compare: A comparator function used to assign priority to elements
   * @returns A priority queue containing the elements from the array
   *
   * @since v0.6.0
   * @history v0.5.4: Originally in `"immutablepriorityqueue"` module with `compare` being a required argument
   */
  provide let fromArray = (array, compare=compare) => {
    Array.reduce((pq, val) => push(val, pq), make(compare=compare), array)
  }
}