K Closest Points to Origin
Você recebe n pontos em um plano: o ponto i é (xs[i], ys[i]). Retorne os k pontos mais próximos da origem (0, 0), cada um como um par [x, y].
Ordene a resposta pela distância da origem, do mais próximo ao mais distante. Pontos à mesma distância vêm primeiro pelo menor x, depois pelo menor y. Um ponto listado mais de uma vez conta uma vez por cópia.
Função
- xsinteger-array
- a coordenada x de cada ponto
- ysinteger-array
- a coordenada y de cada ponto
- kinteger
- quantos pontos retornar
- Retornainteger-2d-array
- os k pontos mais próximos como pares [x, y], em ordem do mais próximo, com os empates desfeitos por x e depois por y
Restrições
n == xs.length == ys.length1 ≤ k ≤ n ≤ 5000-104 ≤ xs[i], ys[i] ≤ 104- Os pontos podem se repetir; cada ocorrência é um ponto separado.
Exemplos
- Entrada
- xs = [3, -1, 2, 0]ys = [1, 2, -2, 5]k = 2
- Saída
- [[-1, 2], [2, -2]]
- Explicação
- As distâncias ao quadrado são 10, 5, 8 e 25. As duas menores pertencem a
(-1, 2), com 5, e a(2, -2), com 8, nessa ordem.
- Entrada
- xs = [1, -1, 0, 1]ys = [1, 1, 2, -1]k = 3
- Saída
- [[-1, 1], [1, -1], [1, 1]]
- Explicação
(1, 1),(-1, 1)e(1, -1)têm todos distância ao quadrado igual a 2, e(0, 2)tem 4. Entre os pontos empatados,(-1, 1)tem o menor x, e(1, -1)vem antes de(1, 1)em relação a y.
- Entrada
- xs = [7]ys = [-7]k = 1
- Saída
- [[7, -7]]
- Explicação
- Há apenas um ponto e você precisa de um, então a resposta é o próprio
(7, -7).
Para ir além
Os pontos agora chegam um de cada vez, e após cada um você deve informar os k mais próximos até então. De quanta memória você precisa?
Dicas
Abra uma de cada vez. Cada uma revela um pouco mais.
Você nunca precisa da distância real. O que pode comparar em vez de
√(x² + y²)e ainda obter a mesma ordem?Digamos que você mantenha os
kpontos mais próximos vistos até agora. Com qual deles um novo ponto compete?Mantenha um max-heap com no máximo
kpontos ordenados por(x² + y², x, y). Substitua o elemento do topo quando um novo ponto vier antes dele e, no final, retire os pontos do heap para preencher a resposta de trás para frente.
Solução
A distância não precisa de raiz quadrada: x² + y² fornece a mesma ordem e continua sendo um número inteiro. Você precisa apenas de k dos n pontos, então um heap máximo dos k pontos mais próximos até o momento, com o mais distante no topo, pronto para ser substituído, custa O(n log k) em vez de uma ordenação completa.
Escolha o ponto mais próximo k vezes
Correta, mas não termina nos maiores testes
Intuição
O primeiro ponto da resposta é o mais próximo de todos, o segundo é o mais próximo dos restantes, e assim por diante. Portanto, percorra os pontos k vezes, escolhendo a cada vez o ponto que vem primeiro segundo (x² + y², x, y) entre os que ainda não foram escolhidos.
Cada percurso verifica todos os n pontos. Para n = 5000 e k = 2500, isso corresponde a 1.25 × 10^7 comparações de três partes.
Algoritmo
- Marque todos os pontos como não selecionados.
- Repita
kvezes: percorra os pontos não selecionados e encontre aquele que vem primeiro segundo(x² + y², x, y). - Marque-o como selecionado e adicione
[x, y]à resposta. - Retorne a resposta.
def kClosest(xs, ys, k):
n = len(xs)
taken = [False] * n
result = []
for _ in range(k):
best = -1
for i in range(n):
# Compare (squared distance, x, y): the order the answer uses.
if not taken[i] and (best == -1 or
(xs[i] * xs[i] + ys[i] * ys[i], xs[i], ys[i]) <
(xs[best] * xs[best] + ys[best] * ys[best], xs[best], ys[best])):
best = i
taken[best] = True
result.append([xs[best], ys[best]])
return resultOrdene cada ponto
Intuição
A resposta são os primeiros k pontos da ordem completa; portanto, ordene todos os pontos por (x² + y², x, y) e mantenha os primeiros k. No exemplo 1, as chaves ordenadas são (5, -1, 2), (8, 2, -2), (10, 3, 1) e (25, 0, 5), e as duas primeiras dão a resposta.
A ordenação custa O(n log n), qualquer que seja o valor de k. Quando k é pequeno, a maior parte desse trabalho ordena pontos que nunca chegam à resposta.
Algoritmo
- Construa a chave de cada ponto
(x² + y², x, y). - Ordene os pontos pela chave.
- Retorne os primeiros
kdeles como pares[x, y].
def kClosest(xs, ys, k):
points = list(zip(xs, ys))
# Sort by squared distance, then x, then y: the order the answer uses.
points.sort(key=lambda p: (p[0] * p[0] + p[1] * p[1], p[0], p[1]))
return [[x, y] for x, y in points[:k]]Mantenha os k elementos mais próximos em um heap máximo
Intuição
Mantenha os melhores k pontos encontrados até agora. Um novo ponto só pode remover o ponto mantido que ocupa a última posição no ranking; por isso, armazene os pontos mantidos em um max-heap ordenado por (x² + y², x, y), com esse ponto no topo. Depois que k pontos estiverem armazenados, um ponto que fique à frente do topo o substitui; qualquer outro ponto é descartado, pois os k pontos armazenados já ficam à frente dele.
No exemplo 2, com k = 3, o topo após os três primeiros pontos é (0, 2), com a chave (4, 0, 2). O ponto (1, -1) tem a chave (2, 1, -1), fica à frente e o substitui.
Um heap só fica parcialmente ordenado, então remova o topo k vezes e preencha a resposta da última posição até a primeira. Cada ponto custa no máximo uma etapa de heap O(log k), então o tempo é O(n log k) e a memória é O(k).
Algoritmo
- Crie um max-heap vazio com prioridade definida por
(x² + y², x, y). - Insira cada um dos primeiros
kpontos. - Para cada ponto posterior, compare-o com o topo: se tiver prioridade maior, substitua o elemento do topo e faça a descida; caso contrário, ignore-o.
- Remova o elemento do topo
kvezes; cada ponto removido ocupa a última posição livre da resposta. - Retorne a resposta.
import heapq
def kClosest(xs, ys, k):
# heapq is a min-heap; negating every part of (distance, x, y) puts the
# farthest kept point, in the answer's order, on top.
heap = []
for x, y in zip(xs, ys):
item = (-(x * x + y * y), -x, -y)
if len(heap) < k:
heapq.heappush(heap, item)
elif item > heap[0]:
# The new point comes before the farthest kept one: swap them.
heapq.heapreplace(heap, item)
# Popping gives the farthest first, so fill the answer from the back.
result = [None] * len(heap)
for i in range(len(result) - 1, -1, -1):
_, x, y = heapq.heappop(heap)
result[i] = [-x, -y]
return result
Armadilhas e casos extremos
- Ignorar o desempate. O exemplo 2 tem três pontos a uma distância de 2; sem comparar
xe depoisy, eles aparecem na ordem de entrada em vez de(-1, 1),(1, -1),(1, 1). - Manter um min-heap de
kpontos. O topo é o ponto mais próximo, então cada remoção descarta um ponto de que a resposta precisa, e você termina com oskpontos mais distantes. - Ler a resposta diretamente do array do heap. Esse array está apenas parcialmente ordenado; remova os pontos um por um primeiro.
- Remover pontos repetidos. Com
(5, 5)listado três vezes ek = 3, a resposta é esse ponto três vezes.
Perguntas frequentes4
Qual é a complexidade de tempo de K Closest Points to Origin?
Ordenar todos os pontos leva O(n log n). A versão com heap tem complexidade O(n log k) para a varredura, mais O(k log k) para ordenar os k pontos restantes, e mantém apenas O(k) pontos na memória.
Por que usar um max-heap, e não um min-heap, para os k pontos mais próximos?
Um novo ponto compete com o mais distante dos k pontos mantidos, então esse ponto precisa ficar no topo. Um heap mínimo só funciona se armazenar todos os n pontos, para que você possa remover os k mais próximos, o que custa O(n) de memória em vez de O(k).
Você precisa da raiz quadrada para comparar distâncias?
Não. Para números que não são negativos, a < b exatamente quando √a < √b, então x² + y² dá a mesma ordem que a distância. Ele também continua sendo um número inteiro: no máximo 2 × 10^8 aqui, dentro do intervalo de 32 bits.
É possível resolver o problema dos K pontos mais próximos da origem em O(n)?
Em média, sim. O Quickselect particiona os pontos em torno de um pivô aleatório pela mesma ordem (x² + y², x, y) e continua apenas no lado que contém a posição k, em O(n) em média. Os k pontos à esquerda dessa posição são a resposta; ordená-los adiciona O(k log k).
Problemas parecidos
Problemas que usam as mesmas ideias. Resolver dois ou três é o que fixa um padrão.
def kClosest(xs, ys, k):
# Escreva o código aquiTeste #1
Teste #2
Teste #3
Argumentos
xs = [3, -1, 2, 0] ys = [1, 2, -2, 5] k = 2
Saída Esperada
[[-1, 2], [2, -2]]