Priority Queue em Ruby

Estive estudando o básico de estrutura de dados e algoritmos novamente durante minhas streams, meus estudos acabaram casando com um vídeo do Augusto Galego sobre otimizar uma solução de um desafio Leetcode utilizando a estrutura de dados Heap.

No vídeo o Augusto resolve usando a própria lib nativa do Python. Em Ruby não existe Heap nativo, porém, existe essa lib de algoritmos do Kanwei. Nela temos as estruturas abstratas e concretas: Priority Queues, Heap (min e max), Stack e etc.

Decidi fazer minha implementação de Priority Queue usando Min Heap/Max Heap, mas antes de sairmos falando sobre a implementação vamos definir o que são essas estruturas que estamos tanto falando, Heap pra lá, Priority Queue pra cá...

Ah, e todo esse processo foi gravado durante uma stream. Se te interessar está aqui os links.

Priority Queue

Uma estrutura de dados abstrata. Quase idêntica ao conceito de Queues (FIFO), porém, a saída de um elemento da Priority Queue respeita uma prioridade contida em cada elemento da fila.

Heap

A estrutura de dados que vamos utilizar na Priority Queue. Ela é uma árvore binária balanceada, e podemos representar ela com um array pra tornar as funções da fila possíveis. Dada a seguinte árvore binária:

Representamos com o seguinte array:
[2, 10, 12, 122, 13, 100]

Nota-se que os números azuis na imagem é o index dos arrays.
E a imagem também nos leva ao próximo tópico sobre Heaps...

Min/Max Heap

Uma heap pode ordenar sua árvore do menor para o maior como na imagem do tópico acima, chamamos de Min Heap. Mas também pode ser possível ordená-la de forma crescente:

Acima, o exemplo de uma Max Heap, que ficaria assim num array
[122, 13, 100, 2, 10, 12]

Inserindo elementos na Heap

Sempre que formos inserir um novo elemento na árvore (ou no array se preferir), temos que rebalancear a árvore caso necessário. Se inserirmos um elemento 6 (com prioridade 6), ela ficaria desbalanceada. O que nos leva a necessidade de compararmos o elemento pai do recém inserido, fazendo um processo chamado de bubble up (ou sift up, heap up:

• Elemento pai maior?
  • Sim? Trocamos de posição o elemento inserido com o pai. (Fazer esse passo até resposta negativa)
  • Não? Finalizamos o processo, temos uma árvore balanceada pós push.

Removendo elementos da heap

Quando removemos um item da heap, vamos rebalancear da mesma forma. Trocamos o elemento da raiz da Heap com o último elemento (index 0 com o último index do array), removemos o último elemento do array e então vamos para o processo de descer o elemento da raíz até o seu devido lugar na árvore chamado de bubble down (ou sift down, heap down):

• Elementos filhos esquerdo ou direito menor?
  • Sim? Priorizar o filho que tiver menor valor contido. Repetir esse passo até resposta negativa.
  • Não? Processo finalizado e temos uma árvore balanceada.

Implementação

Esse com certeza não é o código mais performático de uma Priority Queue, estou compartilhadno e usando tanto este post quanto minha stream como forma de documentar o aprendizado. Pretendo melhorar o código e ele está disponível no repositório de estudos de algorimos no meu github. Isso obviamente não quer dizer que não estou sucinto a comentários sugerindo melhorias, todos são muito bem vindos! :)

Os elementos

Criei uma classe que representaria os elementos da Heap, e nele temos duas importantes coisas a se notar.

class Element
  attr_accessor :data, :priority

  def initialize(data, priority = data)
    @data, @priority = data, priority
  end

  def <=>(elem)
    @priority <=> elem.priority
  end
end

• É possível inicializar um elemento com uma prioridade customizável, mas caso não é inicializada com a prioridade explícita, usaremos o próprio dado como prioridade. Isso vai nos ajudar na criação de uma Heap com inteiros por exemplo: um Element.new(10) já define o elemento 10 com prioridade 10.
• Uma forte característica da Heap é que estaremos comparando o tempo todo os elementos (maior que, menor que). Criamos uma função que sobreescreve o spaceship operator (<=>). O que ele faz é simples:

a <=> b
#  se a < b retorna -1
#  se a = b retorna 0
#  se a > b retorna return  1
#  se a e b não são comparáveis, retorna nil

Quando utilizarmos element1 <=> element2 ele vai nos retornar um desses 4 valores. Esse valor vai determinar a posição de um elemento após ser removido da fila, ou enfileirado. Isso vai nos ajudar na lógica de bubble up ou bubble down

A Priority Queue

Vamos por partes...

class PriorityQueue
  attr_accessor :elements

  def initialize(elements = [], order = -1)
    @elements = elements
    @compare = lambda { |left_el, right_el| (left_el <=> right_el) == order}
  end

  def size
    @elements.size
  end
end

E aqui é onde o spaceship operator entra em ação. A lambda utilizando o order com 1 negativo quer dizer que sempre que compararmos um valor menor com um maior, retornará true. Isso nos garante que ao compararmos o elemento com um pai por exemplo, ele vai retornar true caso precise trocar

Enfileirando...

Ao enfileirar um elemento na Priority Queue, veremos se é necessário o processo de bubble up. A vantagem de traduzirmos a árvore binária para um array é que conseguimos computar facilmente o parente daquele elemento inserido. Existe a fórmula pronta pra conseguir o index do parente, do filho da esquerda e da direita. No caso de enfileiras vamos usar só a fórmula do parente.

PARENT = (TARGET_EL_INDEX - 1) / 2

  def <<(element)
    # Insert new element into array
    @elements << element

    # Get the inserted element index
    target_index = size - 1

    # Get the parent from the inserted index
    parent_index = (target_index - 1) / 2

    # Here, we use compare lambda to make sure we always swap the parent with the inserted element, to respect the min heap property (The child priority must be greater than parent's priority)
    while @compare.call(@elements[target_index], @elements[parent_index])
      # parent greater than child? swap it
      swap(target_index, parent_index)

      # update the respective indexes if next @compare.call succeeds
      target_index = parent_index

      parent_index = (target_index - 1) / 2
    end
  end

Nota-se que a order da PriorityQueue vai seguir a padrão (-1), ou seja, uma Min Heap.

Pop, pop, pop

Para o processo de bubble down, vamos usar as duas fórmulas pra encontrar o filho da esquerda e o filho da direita do elemento que precisa descer na árvore.

LEFT_CHILD = 2 * (TARGET_EL_INDEX) + 1
RIGHT_CHILD = 2 * (TARGET_EL_INDEX) + 2

  def pop
    # swap the binary tree root with the last element
    swap(0, size - 1)

    # remove last element from queue
    @elements.pop

    # computations to reach the left/right child
    target_index = 0
    left_index = 2 * target_index + 1
    right_index = 2 * target_index + 2

    # here we start the bubbling up procedure and we dont stop until our childs elements is greater than its parents
    while (left_index < size && @compare.call(@elements[left_index], @elements[target_index])) || (right_index < @elements.size && @compare.call(@elements[right_index], @elements[target_index]))
      # if right and left child holds the same priority, then we use the data contained into it to decide which is smaller
      smallest_index = if right_index >= size || @elements[left_index].data < @elements[right_index].data
                         left_index
                       else
                         right_index
                       end

      # we swap it boooyyy
      swap(target_index, smallest_index)

      # update indexes in case this was not the last swap we needed it
      target_index = smallest_index
      left_index = 2 * target_index + 1
      right_index = 2 * target_index + 2
    end
  end

Amostras & conclusões

Vamos lá...
Criei um método inspect dentro da class PriorityQueue pra ficar fácil o debugging.

  def inspect
    res = []
    for el in @elements
      res.push(el.data)
    end

    res
  end

E voilà:

irb >>
min_heap = []

min_heap << Element.new(4)
min_heap << Element.new(8)
min_heap << Element.new(6)
min_heap << Element.new(9)
min_heap << Element.new(12)
min_heap << Element.new(7)
min_heap << Element.new(10)

pq = PriorityQueue.new(min_heap)
p pq.inspect()
=> [4, 8, 6, 9, 12, 7, 10]

pq << Element.new(5)
p pq.inspect()
=> [4, 5, 6, 8, 12, 7, 10, 9]

pq.pop()
p pq.inspect()
=> [5, 8, 6, 9, 12, 7, 10]

pq.pop()
p pq.inspect()
=> [6, 8, 7, 9, 12, 10]

pq.pop()
pq.pop()
pq.pop()
p pq.inspect()
=> [9, 12, 10]

Na stream procurei fazer os testes e esse debugging usando excalidraw pra facilitar o entendimento dos resultados.

Pretendo melhorar esse código, colocar novos métodos (por exemplo um pop que tem index como argumento). Se te interessar, pode ficar ligado no meu github e nas streams. Até mais ver. o/