Árboles k-d para búsqueda espacial

Qué cubre este capítulo

Los árboles binarios de búsqueda organizan números sobre una línea para que puedas encontrar uno en O(log n). Pero muchísimos datos viven en el espacio, no en una línea — puntos en un mapa, pixeles en el espacio de color, vectores de features en machine learning — y las preguntas son espaciales: ¿cuál punto guardado está más cerca de este otro? ¿cuáles puntos caen dentro de esta caja? Un árbol k-d generaliza el árbol binario de búsqueda a k dimensiones para responder justo eso, muchísimo más rápido que revisar cada punto. Divide el espacio un eje a la vez — la raíz corta en x, sus hijos en y, los de ellos otra vez en x — así cada nodo parte su región en dos, y una búsqueda del vecino más cercano puede podar subárboles completos que demostrablemente están demasiado lejos como para importar. Este capítulo construye el árbol y su búsqueda con poda, observa cómo una consulta encuentra a su vecino más cercano saltándose la mayoría de los puntos, y encara la famosa limitación de la estructura: conforme crecen las dimensiones, la poda deja de funcionar, hasta que en 16 dimensiones el árbol k-d examina el 99% de los puntos — nada mejor que fuerza bruta. Esa "maldición de la dimensionalidad" es una de las ideas más importantes de toda la ciencia de datos.

Un poco de historia

Jon Bentley presentó los árboles k-d en 1975, cuando era estudiante de posgrado, como una generalización multidimensional del árbol binario de búsqueda — la "k-d" es literalmente "k-dimensional". Fue parte de una ola de trabajo sobre geometría computacional y estructuras de datos espaciales en los setentas y ochentas (los quadtrees, los R-trees y otros llegaron por esas mismas fechas), cuando la computación empezó a lidiar con datos geográficos, gráficos y científicos que eran espaciales por naturaleza. El árbol k-d se volvió la estructura estándar para búsqueda del vecino más cercano en pocas dimensiones y sigue usándose muchísimo (SciPy, scikit-learn y un montón de sistemas de gráficos y robótica lo incluyen). Pero su historia también es la historia de aprender sus límites: conforme los investigadores lo aplicaron a datos de más dimensiones — sobre todo en machine learning, donde los vectores de features tienen decenas o cientos de dimensiones — vieron cómo su eficiencia se venía abajo, y ese fracaso cristalizó el entendimiento de la "maldición de la dimensionalidad", una frase que Richard Bellman había acuñado en 1957 para la explosión general del volumen con la dimensión. El árbol k-d es entonces una doble lección: una estructura elegante que funciona de maravilla en las pocas dimensiones para las que fue diseñada, y una demostración concreta de por qué la geometría de muchas dimensiones derrota nuestras intuiciones de pocas dimensiones.

La intuición

Construir un árbol k-d es como construir un árbol binario de búsqueda, pero la "llave" va rotando entre las dimensiones. En la raíz, separa los puntos por su coordenada x: toma el punto con la x mediana como raíz, manda los puntos con x menor al subárbol izquierdo y los de x mayor al derecho. En el siguiente nivel, separa por y. En el nivel de abajo, de vuelta a x. Cada nodo representa entonces un plano de corte (una línea vertical en 2D en la raíz, luego líneas horizontales, alternando) que divide su región del espacio en dos mitades, y el árbol particiona recursivamente todo el espacio en celdas rectangulares, una por punto. Cortar en la mediana mantiene el árbol balanceado, con profundidad O(log n).

Ahora la recompensa, la búsqueda del vecino más cercano, que es donde vive la poda. Para encontrar el punto más cercano a una consulta, baja por el árbol hacia el lado de cada corte donde cae la consulta — así llegas rapidísimo a la celda chiquita donde está la consulta, y su punto es un buen primer candidato a más cercano. Luego regresa hacia arriba. En cada nodo ya buscaste el lado propio de la consulta; la pregunta es si el otro lado podría tener algo más cerca. La respuesta es una sola comparación: ¿qué tan lejos está la consulta del plano de corte? Si esa distancia perpendicular es mayor que la distancia al mejor punto encontrado hasta ahora, entonces nada del otro lado del plano puede estar más cerca — todo ese subárbol se poda, jamás se visita. Solo si la consulta está lo bastante cerca del plano como para que pudiera esconderse algo más cercano del otro lado tienes que buscar allá. En pocas dimensiones, la mayoría de los subárboles se podan así, de modo que la búsqueda toca O(log n) puntos en lugar de los n. La animación lo enseña: una consulta en una esquina encuentra a su vecino más cercano tras visitar apenas 4 de 12 puntos, porque los planos de corte le permiten demostrar que los otros 8 están demasiado lejos como para molestarse en revisarlos.

Complejidad: cómo escala

Construir el árbol es O(n log n) (un corte en la mediana en cada uno de los log n niveles). Las consultas de vecino más cercano y de rango son O(log n) en promedio en pocas dimensiones. Pero esa salvedad de "pocas dimensiones" es toda la historia del árbol k-d, y el enfrentamiento la mide directo: la fracción de todos los puntos que una búsqueda de vecino más cercano de verdad examina, conforme crece el número de dimensiones:

En 2 dimensiones, el árbol k-d examinó cerca del 0.8% de los 2,000 puntos para encontrar un vecino más cercano — cien veces más rápido que fuerza bruta, exactamente la ganancia O(log n) por la que es famoso. Pero mira cómo trepa la curva: para 8 dimensiones ya examina una buena parte, y para 16 dimensiones examina el 99.2% de los puntos — o sea todos, nada mejor que el recorrido O(n) de fuerza bruta, y encima con el costo extra de recorrer el árbol. La poda dejó de funcionar por completo. Esta es la maldición de la dimensionalidad hecha concreta, y no es un defecto del árbol k-d — es un hecho sobre la geometría de muchas dimensiones del que ninguna estructura espacial exacta se escapa. El árbol k-d es excelente en las pocas dimensiones para las que fue construido e inútil en muchas, y saber dónde ocurre esa transición (a grandes rasgos, cuando la dimensión supera log₂ n) es esencial para usarlo — o para saber que no debes.

En qué es bueno y en qué no

Los árboles k-d son la herramienta correcta para consultas espaciales en pocas dimensiones — digamos de 2 a unas 10 — donde su poda da una aceleración O(log n) de verdad. Los usos clásicos: búsqueda de vecino más cercano y de los k vecinos más cercanos sobre datos geográficos o de mapas (encuentra la tienda, la ciudad o el sensor más cercano), detección de colisiones y ray tracing en gráficos y videojuegos, bases de datos espaciales y sistemas de información geográfica, robótica (planeación de movimiento, procesamiento de nubes de puntos de LIDAR), y clasificación y regresión por vecino más cercano en machine learning con pocas dimensiones (el KNeighbors de scikit-learn usa un árbol k-d o un ball tree para datos de pocas dimensiones). Las consultas por rango o por caja — encuentra todos los puntos de una región — son igual de eficientes. Cuando los datos son genuinamente de pocas dimensiones y más o menos estáticos, el árbol k-d es rápido, exacto y simple.

Donde falla es en muchas dimensiones, como muestra el enfrentamiento sin piedad: pasando las ~10 a 20 dimensiones degenera a una fuerza bruta lenta, así que para los vectores de features de muchas dimensiones comunes en machine learning (embeddings de imágenes, vectores de palabras, cientos de dimensiones) es la estructura equivocada — usa métodos aproximados (LSH, HNSW) o reduce primero la dimensión. También funciona mejor con conjuntos de puntos estáticos: igual que un BST balanceado, no se rebalancea al insertar, así que una racha de inserciones lo puede desbalancear (reconstrúyelo cada tanto, o usa una variante dinámica). Las distancias no euclidianas y los puntos que se actualizan seguido lo ponen en aprietos. Y para conjuntos de puntos muy chicos, la fuerza bruta es más simple e igual de rápida. El punto dulce del árbol k-d es preciso: datos espaciales de pocas dimensiones y más o menos estáticos, donde su elegante poda brilla — y reconocer cuándo te saliste de ese punto dulce (la cantidad de dimensiones) es tan importante como la estructura misma.

Los datos, o las entradas

El enfrentamiento mide la fracción de puntos que un árbol k-d examina para la búsqueda del vecino más cercano conforme la dimensión crece de 2 a 16, sobre 2,000 puntos aleatorios — la maldición de la dimensionalidad en una sola curva. La correctitud se verifica contra fuerza bruta en cientos de conjuntos de puntos aleatorios de 1 a 4 dimensiones: el vecino más cercano del árbol k-d debe estar exactamente a la misma distancia que el de la fuerza bruta (comparando distancias, ya que los empates permiten puntos distintos), y su búsqueda por rango o caja debe devolver exactamente los puntos que encuentra un filtro por fuerza bruta. La animación corre una búsqueda de vecino más cercano en 2D: doce puntos particionados por las líneas de corte del árbol, y una consulta en la esquina superior derecha que encuentra a su vecino más cercano podando la mayor parte del árbol.

Constrúyelo, una función a la vez

Construir el árbol — corte en la mediana con un eje que rota en cada nivel:

def build_kdtree(points, depth=0):
    """Build a balanced k-d tree by, at each level, splitting the points along the current axis at
    their MEDIAN (which keeps the tree balanced, O(log n) deep). The axis cycles through the
    dimensions with depth: 0, 1, …, k-1, 0, 1, …. Left subtree holds points below the median on this
    axis, right subtree holds those above. O(n log n) to build (with a median sort per level)."""
    if not points:
        return None
    k = len(points[0])
    axis = depth % k
    points = sorted(points, key=lambda p: p[axis])
    mid = len(points) // 2
    node = KDNode(points[mid], axis)
    node.left = build_kdtree(points[:mid], depth + 1)
    node.right = build_kdtree(points[mid + 1:], depth + 1)
    return node

Búsqueda del vecino más cercano — baja, y luego poda los subárboles que no pueden estar más cerca:

def nearest(root, target, counter=None):
    """The nearest point in the tree to `target`, with pruning. Descend toward the target's side of
    each split first (it's the likely place the nearest neighbor lives); on the way back up, only
    search the FAR side of a split if the distance from the target to the splitting plane is less than
    the best distance found so far — otherwise nothing on that side can be closer, so the whole
    subtree is pruned. `counter` tallies nodes visited, to compare against brute force. Returns the
    nearest point."""
    best = [None, math.inf]                        # [point, squared distance]

    def visit(node):
        if node is None:
            return
        if counter is not None:
            counter[0] += 1
        d2 = sum((a - b) ** 2 for a, b in zip(node.point, target))
        if d2 < best[1]:
            best[0], best[1] = node.point, d2
        axis = node.axis
        diff = target[axis] - node.point[axis]     # signed distance to the splitting plane
        near, far = (node.left, node.right) if diff < 0 else (node.right, node.left)
        visit(near)                                 # explore the target's own side first
        if diff * diff < best[1]:                   # the far side could hold something closer → search it
            visit(far)                              # otherwise it's pruned entirely

    visit(root)
    return best[0]

Búsqueda por rango — la misma idea de poda para "¿qué puntos están en esta caja?":

def range_search(root, low, high):
    """All stored points inside the axis-aligned box [low, high] (per-dimension bounds). At each node,
    only descend into a child whose region overlaps the box — the same pruning idea applied to a
    region query. Returns the list of points in the box."""
    found = []

    def visit(node):
        if node is None:
            return
        p = node.point
        if all(low[i] <= p[i] <= high[i] for i in range(len(p))):
            found.append(p)
        axis = node.axis
        if low[axis] <= p[axis]:                    # left region might overlap the box
            visit(node.left)
        if p[axis] <= high[axis]:                   # right region might overlap the box
            visit(node.right)

    visit(root)
    return found

Míralo funcionar

Aquí hay un árbol k-d 2D de doce puntos, con los cortes recursivos del árbol dibujados como líneas — verticales para los cortes en x, horizontales para los cortes en y, tallando el plano en celdas. El triángulo rojo es la consulta, arriba a la derecha. Observa la búsqueda del vecino más cercano: baja hacia la esquina de la consulta y llega rápido a puntos cercanos (azul, visitados; verde, el mejor actual), afinando su estimación del más cercano. Los momentos clave son las podas: cuando la consulta está más lejos de una línea de corte que de su mejor punto hasta ese momento, toda la región del otro lado de esa línea se salta — esos puntos no pueden estar más cerca, así que la búsqueda ni siquiera los mira. Encuentra el punto más cercano, (9, 9), habiendo visitado solo 4 de los 12 puntos; los planos de corte le permiten demostrar que los otros 8 están demasiado lejos sin medirlos. Esa demostración por geometría es todo el poder del árbol k-d — y es exactamente lo que destruye la maldición de la dimensionalidad en muchas dimensiones, donde los planos siempre están demasiado cerca para podar:

El código completo

La pestaña "desde cero" es el árbol k-d — construcción, vecino más cercano con poda y búsqueda por rango; la pestaña de librería es el vecino más cercano por fuerza bruta O(n) contra el que se verifica y compite, con una nota sobre scipy.spatial.KDTree para producción. Cámbiate entre ellas — la fuerza bruta es trivialmente correcta y siempre O(n); toda la astucia del árbol k-d (y su caída en muchas dimensiones) está en la prueba de poda.

"""k-d tree — a binary search tree generalized to k dimensions, for organizing points in space so
that spatial queries ("which stored point is nearest to this one?", "which points lie in this
region?") run far faster than checking every point. A binary search tree splits a line of numbers;
a k-d tree splits k-dimensional space, one axis at a time, alternating axes as it descends: the root
splits on the x-coordinate, its children on y, their children on x again, and so on. Each node owns a
point and divides its region of space in two along the current axis.

The payoff is PRUNING. To find the nearest neighbor of a query point, you descend to the leaf region
the query falls in (fast, like a BST search), then walk back up — but at each node you only need to
explore the *other* side of the split if it could possibly contain something closer than the best
point found so far, which you check with a single comparison against the splitting plane. Most of the
tree gets pruned, so nearest-neighbor search is O(log n) on average in low dimensions instead of the
O(n) of checking every point. This chapter builds the tree, its nearest-neighbor search, and reveals
the catch that limits it: as dimensions grow, the pruning stops working — the "curse of dimensionality."
"""
import math


class KDNode:
    __slots__ = ("point", "axis", "left", "right")

    def __init__(self, point, axis):
        self.point = point
        self.axis = axis
        self.left = None
        self.right = None


# region: build
def build_kdtree(points, depth=0):
    """Build a balanced k-d tree by, at each level, splitting the points along the current axis at
    their MEDIAN (which keeps the tree balanced, O(log n) deep). The axis cycles through the
    dimensions with depth: 0, 1, …, k-1, 0, 1, …. Left subtree holds points below the median on this
    axis, right subtree holds those above. O(n log n) to build (with a median sort per level)."""
    if not points:
        return None
    k = len(points[0])
    axis = depth % k
    points = sorted(points, key=lambda p: p[axis])
    mid = len(points) // 2
    node = KDNode(points[mid], axis)
    node.left = build_kdtree(points[:mid], depth + 1)
    node.right = build_kdtree(points[mid + 1:], depth + 1)
    return node
# endregion


# region: nearest
def nearest(root, target, counter=None):
    """The nearest point in the tree to `target`, with pruning. Descend toward the target's side of
    each split first (it's the likely place the nearest neighbor lives); on the way back up, only
    search the FAR side of a split if the distance from the target to the splitting plane is less than
    the best distance found so far — otherwise nothing on that side can be closer, so the whole
    subtree is pruned. `counter` tallies nodes visited, to compare against brute force. Returns the
    nearest point."""
    best = [None, math.inf]                        # [point, squared distance]

    def visit(node):
        if node is None:
            return
        if counter is not None:
            counter[0] += 1
        d2 = sum((a - b) ** 2 for a, b in zip(node.point, target))
        if d2 < best[1]:
            best[0], best[1] = node.point, d2
        axis = node.axis
        diff = target[axis] - node.point[axis]     # signed distance to the splitting plane
        near, far = (node.left, node.right) if diff < 0 else (node.right, node.left)
        visit(near)                                 # explore the target's own side first
        if diff * diff < best[1]:                   # the far side could hold something closer → search it
            visit(far)                              # otherwise it's pruned entirely

    visit(root)
    return best[0]
# endregion


# region: range_search
def range_search(root, low, high):
    """All stored points inside the axis-aligned box [low, high] (per-dimension bounds). At each node,
    only descend into a child whose region overlaps the box — the same pruning idea applied to a
    region query. Returns the list of points in the box."""
    found = []

    def visit(node):
        if node is None:
            return
        p = node.point
        if all(low[i] <= p[i] <= high[i] for i in range(len(p))):
            found.append(p)
        axis = node.axis
        if low[axis] <= p[axis]:                    # left region might overlap the box
            visit(node.left)
        if p[axis] <= high[axis]:                   # right region might overlap the box
            visit(node.right)

    visit(root)
    return found
# endregion

"""The reference and the contrast. There's no k-d tree in the Python standard library; in production
you'd use SciPy:

    from scipy.spatial import KDTree
    tree = KDTree(points)
    dist, idx = tree.query(target)        # nearest neighbor
    idxs = tree.query_ball_point(target, r)   # all points within radius r

`brute_force_nearest` below checks every point — O(n) per query, the baseline a k-d tree beats in low
dimensions (and, tellingly, ties in high dimensions). It's both the correctness reference and the
contrast that reveals the curse of dimensionality.
"""


# region: brute
def brute_force_nearest(points, target):
    """The obvious O(n) nearest-neighbor search: compute the distance to every point and keep the
    closest. Always correct, never pruned — the reference the k-d tree is checked against, and the
    cost it improves on when the dimension is low enough for pruning to work."""
    best, best_d2 = None, float("inf")
    for p in points:
        d2 = sum((a - b) ** 2 for a, b in zip(p, target))
        if d2 < best_d2:
            best, best_d2 = p, d2
    return best
# endregion

Desde cero vs librería

El árbol k-d enseña dos cosas al mismo tiempo. Primero, la lección constructiva: particionar el espacio con una prueba de poda convierte la búsqueda del vecino más cercano de O(n) a O(log n), usando geometría para demostrar que la mayoría de los puntos no pueden ser la respuesta — el mismo espíritu de "elimina lo que no puede importar" que la poda del backtracking y de branch-and-bound. Segundo, y más profunda, la lección de advertencia: esa hermosa aceleración se evapora en muchas dimensiones, y el ascenso del enfrentamiento de 0.8% a 99% es una de las demostraciones más claras de la maldición de la dimensionalidad que puedes correr. Esa maldición no es un bug del árbol k-d — es una propiedad fundamental del espacio de muchas dimensiones que derrota a todo método exacto, y es por eso que la búsqueda moderna de vecino más cercano en muchas dimensiones (las bases de datos vectoriales detrás de la búsqueda semántica y RAG, la recuperación en los sistemas de recomendación) usa algoritmos aproximados (HNSW, LSH). Entender el árbol k-d — tanto su poder en pocas dimensiones como su colapso en muchas — es entender por qué existe toda esa industria. En producción usarías el KDTree de scipy para datos de pocas dimensiones y una base de datos vectorial (FAISS, basada en HNSW) para los de muchas; construir el árbol k-d tú mismo es lo que hace que tanto la poda como su desplome sean algo que mediste en vez de algo que te advirtieron.

Dónde te lo vas a encontrar

Los árboles k-d corren donde sea que aparezcan consultas espaciales de pocas dimensiones. Los servicios de mapas y ubicación los usan (y a sus parientes) para consultas de "la gasolinera / el conductor / el amigo más cercano". Los gráficos por computadora y los videojuegos los usan para ray tracing, detección de colisiones y photon mapping. La robótica los usa para planeación de movimiento y para procesar nubes de puntos de LIDAR (3D). El machine learning los usa para clasificación y regresión por k vecinos más cercanos sobre features de pocas dimensiones (la implementación de scikit-learn) y para estimación de densidad. El cómputo científico y la simulación los usan para interacciones de partículas (problemas de N cuerpos, hidrodinámica de partículas suavizadas) y estadística espacial. Las bases de datos usan a sus primos de índices espaciales (los R-trees) para consultas geográficas. Y algo clave: la ausencia de árboles k-d en la búsqueda vectorial de muchas dimensiones — reemplazados por HNSW y LSH en FAISS, Pinecone, Milvus y toda base de datos vectorial — es consecuencia directa de la maldición que demuestra este capítulo. Donde haya que buscar puntos en un espacio de pocas dimensiones, lo más probable es que la herramienta sea un árbol k-d o un pariente cercano; donde sean de muchas dimensiones, su fracaso es la razón de que sea otra cosa.

Puntos clave

Un árbol k-d generaliza el árbol binario de búsqueda a k dimensiones, cortando el espacio con un eje que rota en cada nivel para que las consultas de vecino más cercano y de rango puedan podar subárboles que demostrablemente no pueden contener un punto más cercano — O(log n) en pocas dimensiones, cien veces más rápido que fuerza bruta en 2D. Pero esa poda depende de que el vecino más cercano esté de verdad cerca en relación con los planos de corte, cosa que falla en muchas dimensiones: para 16 dimensiones el árbol k-d examina el 99% de los puntos, nada mejor que fuerza bruta. Ese colapso es la maldición de la dimensionalidad — una propiedad fundamental de la geometría de muchas dimensiones, no un defecto de la estructura — y es por eso que los métodos espaciales exactos le ceden el lugar a los aproximados (HNSW, LSH) para los datos de muchas dimensiones del machine learning moderno.

El capítulo final del libro es una joyita perfecta que amarra todo el nivel avanzado. El reservoir sampling saca una muestra aleatoria uniforme de un stream de longitud desconocida y posiblemente infinita — datos demasiado grandes para guardarlos, vistos una sola vez — usando memoria fija y una sola pasada. Es la respuesta a "cómo muestreas de forma justa algo que no puedes sostener", una pregunta que las estructuras anteriores ni tocan, y cierra el libro con el tema recurrente del nivel: que la idea correcta, muchas veces una probabilística, vuelve rutinario lo que parecía imposible.