JavaScript

JavaScript、Golang 手写实现优先队列，PHP 重载 SplPriorityQueue 的 compare 方法实现最小堆，求解《23. 合并K个升序链表》和《剑指 Offer II 078. 合并排序链表》

代码模板 · 优先队列

class PriorityQueue {
  constructor(ar, compare = (a, b) => a - b) {
    this.heap = []
    this.size = 0
    this.compare = compare
    while (ar.length) this.push(ar.pop())
  }
  swap(a, b) {
    const t = this.heap[a]
    this.heap[a] = this.heap[b]
    this.heap[b] = t
  }
  getLeftIndex(index) {
    return (index << 1) + 1
  }
  getRightIndex(index) {
    return (index << 1) + 2
  }
  getParentIndex(index) {
    return index - 1 >> 1
  }
  shiftDown(index) {
    const leftIndex = this.getLeftIndex(index)
    const rightIndex = this.getRightIndex(index)
    if (this.heap[leftIndex] !== void 0 && this.compare(this.heap[index], this.heap[leftIndex]) > 0) {
      this.swap(leftIndex, index)
      this.shiftDown(leftIndex)
    }
    if (this.heap[rightIndex] !== void 0 && this.compare(this.heap[index], this.heap[rightIndex]) > 0) {
      this.swap(rightIndex, index)
      this.shiftDown(rightIndex)
    }
  }
  shiftUp(index) {
    const parentIndex = this.getParentIndex(index)
    if (this.heap[parentIndex] !== void 0 && this.compare(this.heap[parentIndex], this.heap[index]) > 0) {
      this.swap(parentIndex, index)
      this.shiftUp(parentIndex)
    }
  }
  push(value) {
    this.heap.push(value)
    this.shiftUp(this.size++)
  }
  top() {
    return this.heap[0]
  }
  pop() {
    if (this.size === 0) return null
    if (--this.size === 0) return this.heap.pop()
    const top = this.top()
    this.heap[0] = this.heap.pop()
    this.shiftDown(0)
    return top
  }
  isEmpty() {
    return this.size === 0
  }
}

type PriorityQueue struct {
  heap []*ListNode
  size int
  compare func(l1 *ListNode, l2 *ListNode) int
}
func (pq *PriorityQueue) Swap(a int, b int) {
  pq.heap[a], pq.heap[b] = pq.heap[b], pq.heap[a]
}
func (pq *PriorityQueue) GetParentIndex(index int) int {
  return (index - 1) >> 1
}
func (pq *PriorityQueue) ShiftUp(index int) {
  parentIndex := pq.GetParentIndex(index)
  if parentIndex >= 0 && pq.compare(pq.heap[parentIndex], pq.heap[index]) > 0 {
    pq.Swap(parentIndex, index)
    pq.ShiftUp(parentIndex)
  }
}
func (pq *PriorityQueue) GetLeftIndex(index int) int {
  return (index << 1) + 1
}
func (pq *PriorityQueue) GetRightIndex(index int) int {
  return (index << 1) + 2
}
func (pq *PriorityQueue) ShiftDown(index int) {
  leftIndex, rightIndex := pq.GetLeftIndex(index), pq.GetRightIndex(index)
  if leftIndex < pq.size && pq.compare(pq.heap[index], pq.heap[leftIndex]) > 0 {
    pq.Swap(leftIndex, index)
    pq.ShiftDown(leftIndex)
  }
  if rightIndex < pq.size && pq.compare(pq.heap[index], pq.heap[rightIndex]) > 0 {
    pq.Swap(rightIndex, index)
    pq.ShiftDown(rightIndex)
  }
}
func (pq *PriorityQueue) Push(value *ListNode) {
  pq.heap = append(pq.heap, value)
  pq.ShiftUp(pq.size)
  pq.size++
}
func (pq *PriorityQueue) Top() *ListNode {
  return pq.heap[0]
}
func HeapPop(heap []*ListNode) (last *ListNode, r []*ListNode) {
  n := len(heap)
  last = heap[n - 1]
  r = heap[:n - 1]
  return
}
func (pq *PriorityQueue) Pop() *ListNode {
  if pq.size == 0 {
    return nil
  }
  pq.size--
  if pq.size == 0 {
    last, heap := HeapPop(pq.heap)
    pq.heap = heap
    return last
  }
  top := pq.Top()
  last, heap := HeapPop(pq.heap)
  pq.heap = heap
  pq.heap[0] = last
  pq.ShiftDown(0)
  return top
}
func (pq *PriorityQueue) IsEmpty() bool {
  return pq.size == 0
}

// 最大堆（默认）
class PriorityQueue extends SplPriorityQueue {
  public function compare($priority1, $priority2) {
    return $priority1 - $priority2;
  }
}
// 最小堆
class PriorityQueue extends SplPriorityQueue {
  public function compare($priority1, $priority2) {
    return $priority2 - $priority1;
  }
}

例题

答案

优先队列

• 将多个链表头部放入优先队列
  • 每次取出最小值，再将头部去除，放入优先队列
  • 直到优先队列为空

var mergeKLists = function(lists) {
  const priorityQueue = new PriorityQueue(lists.filter(v => v), (a, b) => a.val - b.val)
  let r = p = new ListNode()
  while (priorityQueue.isEmpty() === false) { 
    p.next = priorityQueue.pop()
    if (p.next.next) priorityQueue.push(p.next.next)
    p = p.next
  }
  return r.next
};

func mergeKLists(lists []*ListNode) *ListNode {
  priorityQueue, p := &PriorityQueue{
    compare: func(l1 *ListNode, l2 *ListNode) int {
      return l1.Val - l2.Val
    },
  }, &ListNode{}
  r := p
  for _, list := range lists {
    if list != nil {
      priorityQueue.Push(list)
    }
  }
  for priorityQueue.IsEmpty() == false {
    p.Next = priorityQueue.Pop()
    if p.Next.Next != nil {
      priorityQueue.Push(p.Next.Next)
    }
    p = p.Next
  }
  return r.Next
}

class PriorityQueue extends SplPriorityQueue {
  public function compare($priority1, $priority2) {
    return $priority2 - $priority1;
  }
}
class Solution {
  function mergeKLists($lists) {
    $r = $p = new ListNode();
    $priorityQueue = new PriorityQueue();
    $priorityQueue->setExtractFlags(SplPriorityQueue::EXTR_DATA);
    foreach ($lists as $list) {
      if ($list) $priorityQueue->insert($list, $list->val);
    }
    while ($priorityQueue->valid()) {
      $p->next = $priorityQueue->top();
      if ($p->next->next) $priorityQueue->insert($p->next->next, $p->next->next->val);
      $p = $p->next;
      $priorityQueue->next();
    }
    return $r->next;
  }
}