Resolvendo o "Five sort"

O objetivo desse post é demonstrar como eu pensei e solucionei o desafio citado usando uma técnica chamada two pointers.

Desafio

Dado um array de números inteiros, mova os números 5 para o final dele. Exemplo: [1, 5, 2, 5, 3]
Como resultado, espera-se o seguinte retorno da função: [1, 3, 2, 5, 5] ou [2, 1, 3, 5, 5]. O importante é que os 5 estejam no final do array.

Abordagem

Há várias maneiras de resolver esse problema. Muitas pessoas deram ótimas soluções nessa thread no Twitter.

Vou abordar nesse texto uma forma mais profunda de ver o problema além do código.

Uma das maneiras de abordar esse problema é removendo todos os 5 da lista original, adicioná-los em outros lista e concatenar no final (lista1 + lista2). Ou remover da lista e adicionar no final.

Não é uma solução ruim, porém, ela tem uma alta complexidade de tempo. Toda vez que você remove um item de um array e ele não é o último, o array precisa ser rearranjado. Isso significa que toda remoção vai adicionar N passos adicionais na sua solução.

Uma abordagem interessante pro problema é usar a técnica two pointers, assumindo que podemos modificar o array in-place, ou seja, modificar o valor do array que, em Python, por exemplo, é mutável.

A técnica se baseia em ter dois ponteiros para fazer comparações em diferentes partes do array, enquanto se move e aplica todas as operações necessárias para resolver o problema.

Nossa abordagem irá comparar os elementos do início com os do fim do array. Com um exemplo de dois elementos ([5, 1]) fica mais fácil. Começando o algoritmo, comparamos o primeiro valor com o último. O primeiro é 5, então movemos ele pro fim, trocando pelo 1, resultando em [1, 5].

Vamos adicionar dois valores [5, 5, 1, 3].
Primeiro e últimos trocam, pois o primeiro é 5. Movemos os dois ponteiros, trocando 5 e 1 também, resultando em [3, 1, 5, 5].

Abaixo uma demonstração visual (com texto alternativo) do problema sendo resolvido:

Nosso algoritmo se baseia em comparar os elementos do array em direção ao meio do dele. Enquanto o ponteiro i for menor que o ponteiro j, nós continuamos a comparar os elementos e trocando quando for necessário.

Código

A nossa implementação vai usar Python, uma linguagem muito utilizada em resolução de algoritmos, tanto em plataformas quanto leetcode quando em entrevistas de emprego.

def five_sort(nums):
  i, j = 0, len(nums) - 1

  while i < j:
    if nums[j] == 5:
      j -= 1
      continue

    if nums[i] == 5:
      nums[j], nums[i] = 5, nums[j]
      j -= 1

    i += 1

  return nums

Passos do algoritmo:

• Começamos os ponteiros nas pontas (começo e fim);
• Respeitamos nosso caso base de iterar até que os ponteiros não tenham se encontrado;
• Quando o 5 já está no fim do array, apenas decrementamos o j para que continuemos a avaliar os outros números;
• Quando o 5 é encontrado no ponteiro i, precisamos fazer a troca com o número do ponteiro j;
• Retornamos a lista modificada in-place.

Complexidade

Tempo: O(n)

Independente de quantos elementos tem o array, no pior caso precisamos iterar em todos elementos.

Espaço: O(1)

Nenhum espaço adicional é alocado. Como fazemos a modificação in-place, não há a criação de um array adicional.

Comments

[Eduardo Klosowski]

Estava olhando a função apresentada, e em toda repetição do laço é validado se nums[j] não é 5, o que lida com os 5 que estão no final da lista, mas é um acesso à memória que o algoritmo precisa fazer. Pensando em remover esse primeiro if, é possível tratar nums[j] como incerto se é ou não 5. Embora contra intuitivo, eu acredito que abrir mão da garantia que esse if da seja uma boa opção também. Como resultado, o algoritmo fica mais simples de ler por ter menos caminhos possíveis dentro do laço de repetição, é fácil identificar o que é feito em cada caso, e a cada execução a garantia de mover apenas um dos ponteiros uma única posição:

def five_sort(nums):
    i, j = 0, len(nums) - 1

    while i < j:
        if nums[i] == 5:
            nums[j], nums[i] = 5, nums[j]
            j -= 1
        else:
            i += 1

    return nums

A desvantagem é que pode causar um pouco mais de trocas dos valores na memória. O que você acha?

[Maurício Antunes]

De fato ficou mais legível, mesmo, pelo motivo mencionado por você, de ter menos caminhos pra percorrer/assimilar. Obrigado pela contribuição!
Sobre a parte de ter mais swaps, não acho que seja problema.