如何为双列表实现的Scala Queue添加O(1)复杂度的extractMax方法

实现思路 #

我们可以基于原双栈队列的结构，额外维护两个对应in、out栈的最大值辅助栈，实现O(1)复杂度获取最大值：

• 每个辅助栈和对应业务栈的长度完全一致，栈中每个位置存储从当前位置到业务栈栈底的最大值
• 入队、出队操作时同步更新辅助栈，均摊时间复杂度仍为O(1)
• 队列的全局最大值就是两个辅助栈栈顶的较大值，可直接在O(1)时间返回
• 通过Scala的上下文界定A: Ordering约束泛型可比较，天然支持Integer/Int类型的需求

完整实现代码 #

import scala.math.Ordering

class Queue[A: Ordering] private (
  private val in: List[A], 
  private val inMax: List[A],
  private val out: List[A],
  private val outMax: List[A]
) {
  private val ord = implicitly[Ordering[A]]

  def isEmpty: Boolean = in.isEmpty && out.isEmpty

  def enqueue(elem: A): Queue[A] = {
    // 计算新的入栈最大值栈头：当前元素和原入栈最大值的较大者
    val newInMaxHead = if (inMax.isEmpty) elem else ord.max(elem, inMax.head)
    new Queue(elem :: in, newInMaxHead :: inMax, out, outMax)
  }

  def dequeue(): (A, Queue[A]) = {
    out match {
      case head :: tail =>
        // 出栈非空直接弹出，同步弹出最大值栈的头
        (head, new Queue(in, inMax, tail, outMax.tail))
      case Nil =>
        if (in.isEmpty) throw new NoSuchElementException("Queue is empty")
        // 入栈倒序转为出栈，同时生成出栈对应的最大值栈
        val reversedIn = in.reverse
        val newOutMax = reversedIn.tail.scanLeft(reversedIn.head)(ord.max).reverse
        (reversedIn.head, new Queue(Nil, Nil, reversedIn.tail, newOutMax.tail))
    }
  }

  def extractMax(): A = {
    if (isEmpty) throw new NoSuchElementException("Queue is empty")
    (inMax.headOption, outMax.headOption) match {
      case (Some(inM), None) => inM
      case (None, Some(outM)) => outM
      case (Some(inM), Some(outM)) => ord.max(inM, outM)
    }
  }
}

// 伴生对象提供初始化方法
object Queue {
  def empty[A: Ordering]: Queue[A] = new Queue[A](Nil, Nil, Nil, Nil)
  def apply[A: Ordering](elems: A*): Queue[A] = elems.foldLeft(empty[A])(_.enqueue(_))
}

逻辑说明 #

• 入队时直接计算当前入栈的最大值，追加到inMax头部即可
• 当out栈清空需要倒转in栈时，通过scanLeft快速生成out对应的最大值栈，这个操作的均摊成本仍然是O(1)
• extractMax仅需比较两个辅助栈的栈顶元素，时间复杂度严格O(1)
• 如果仅需要支持Integer类型，直接把泛型A替换为Integer即可，无需修改核心逻辑

内容的提问来源于stack exchange，提问作者atommaks