Segment trees

De qué trata este capítulo

Aquí hay un problema que un array simple no resuelve bien: tienes una lista de números y necesitas preguntar una y otra vez "¿cuál es la suma de todo lo que está entre la posición i y la j?", mientras además cambias elementos individuales. Si guardas los valores en crudo, cada consulta de rango te cuesta O(n). Si precalculas prefix sums, las consultas bajan a O(1) — pero ahora cada update obliga a reconstruir todo en O(n). Te quedas atorado eligiendo cuál operación va a ser la lenta. El segment tree se niega a elegir: es un árbol de agregados por rango que deja ambas cosas, la consulta de rango y el update puntual, en O(log n). En este capítulo lo construimos, vemos cómo una consulta se rompe en O(log n) piezas precalculadas, y lo medimos ganándole a los dos enfoques con array en una carga de trabajo mixta.

Un poco de historia

Los segment trees vienen de la geometría computacional. En los años setenta, mientras los investigadores armaban algoritmos para problemas del estilo "cuáles de estos segmentos de línea contienen un punto dado", Jon Bentley y otros desarrollaron estructuras de árbol sobre rangos de un eje de coordenadas: el segment tree y sus parientes (interval trees, range trees). La idea central, descomponer cualquier intervalo en una cantidad logarítmica de piezas canónicas precalculadas, resultó ser enormemente general y aplicable a cualquier agregado asociativo sobre rangos, no solo a geometría. Tuvo una segunda vida en la programación competitiva de los noventa en adelante, donde "consulta de rango con updates" es pan de todos los días, y el segment tree (con su extensión de lazy propagation) se volvió la artillería pesada estándar. Más que un solo algoritmo es una plantilla: eliges una operación asociativa — suma, mínimo, máximo, gcd — y el mismo árbol te responde consultas de rango y updates puntuales en O(log n).

La intuición

Arma un árbol binario sobre las posiciones del array. La raíz cubre todo el array y guarda el agregado — digamos la suma — de todo. Sus dos hijos cubren la mitad izquierda y la derecha y guardan sus sumas; los hijos de esos, los cuartos; y así hasta las hojas, que son los elementos individuales. Cada nodo guarda el agregado de un rango contiguo, y el valor de un padre no es más que la combinación de los valores de sus dos hijos.

Ahora la consulta de rango. Para sumar las posiciones [i, j], arrancas en la raíz y miras su rango. Si el rango del nodo queda completamente fuera de [i, j], no aporta nada. Si queda completamente dentro de [i, j], tomas directo su agregado guardado — no bajas, porque la suma de todo ese rango ya está calculada y ahí está. Solo cuando el nodo se traslapa parcialmente te partes en sus dos hijos y recurses. Lo mágico es que cualquier rango, sin importar cómo esté posicionado, queda cubierto por apenas O(log n) de esos nodos "completamente dentro" — los segmentos canónicos — así que la consulta arma su respuesta a partir de un número logarítmico de piezas precalculadas. Los updates son el espejo: cambias una hoja y luego regresas hacia arriba recalculando el agregado de cada ancestro a partir de sus hijos, tocando solo los O(log n) nodos de ese único camino. Ninguna de las dos operaciones toca jamás más que una rebanada logarítmica del árbol.

Complejidad: cómo escala

Construir el árbol es $O(n)$ — cada uno de los ~2n nodos se calcula una sola vez. Una consulta de rango es $O(\log n)$: en cada nivel del árbol la consulta toca cuando mucho un número constante de nodos (la descomposición canónica agrega a lo más dos por nivel), así que el total es proporcional a la altura. Un update puntual también es $O(\log n)$ — un solo camino raíz-a-hoja recalculado. El espacio es $O(n)$. La gráfica corre una carga de 5000 operaciones repartidas mitad y mitad entre consultas de rango y updates, contra los dos enfoques con array:

Con 40000 elementos el segment tree corrió la carga mixta en unos 25 ms, contra los 152 ms del array en crudo (6× más lento, porque la mitad de sus operaciones eran consultas de rango O(n)) y los 2544 ms del array de prefix sums (100× más lento, porque la mitad de sus operaciones eran reconstrucciones O(n)). Cada enfoque con array es rápido en una operación y catastrófico en la otra; el segment tree es O(log n) en ambas, así que gana en cuanto la carga las mezcla — que es justo lo que hacen las cargas reales.

En qué es bueno y en qué no

El segment tree es la estructura correcta cuando necesitas agregados de rango sobre datos que cambian: muchas consultas y muchos updates, intercalados. Maneja cualquier operación asociativa (suma, mínimo, máximo, gcd y más) con el mismo código, y con la extensión de lazy propagation hasta puede aplicar updates a un rango completo en O(log n), no solo a un punto. Esa flexibilidad lo vuelve la opción obligada para problemas de rangos en programación competitiva y una buena alternativa para analítica sobre series de tiempo o arrays mutables.

Sus costos son la complejidad y el peso. Es más código y más memoria (más o menos 2n–4n nodos) que las alternativas, así que si tus datos son estáticos — sin updates — las prefix sums te dan consultas O(1) y deberías usar eso y ya. Y si nada más vas a necesitar sumas de prefijo (no mínimos ni máximos), el Fenwick tree del siguiente capítulo hace el mismo trabajo con una fracción del código y de la memoria. El segment tree justifica su overhead cuando de verdad necesitas datos mutables y consultas de rango de una operación que va más allá de sumas; para los casos más acotados, ganan las herramientas más ligeras.

Los datos, o las entradas

El face-off corre una mezcla 50/50 de consultas de rango y updates puntuales, la carga donde cada enfoque con array cae en su operación lenta la mitad del tiempo y donde brilla el O(log n) parejo del segment tree. La animación construye un segment tree sobre ocho elementos y traza una consulta de rango, para que veas cómo descompone el rango en sus segmentos canónicos.

Constrúyelo, una función a la vez

La consulta es el corazón de todo: la decisión de tres vías (fuera, dentro, parcial) que arma la respuesta a partir de segmentos canónicos:

def query(self, ql, qr, probe=None):
    """Sum of data[ql..qr] in O(log n). At each node: if its range is entirely
    OUTSIDE the query, it contributes 0; if entirely INSIDE, it contributes its
    stored aggregate directly — no need to descend; if it PARTIALLY overlaps, split
    into its two children. Only O(log n) nodes ever land 'entirely inside', and their
    ranges are the canonical pieces the answer is assembled from."""
    return self._query(1, 0, self._n - 1, ql, qr, probe)

def _query(self, node, lo, hi, ql, qr, probe):
    if qr < lo or hi < ql:                       # entirely outside
        return 0
    if ql <= lo and hi <= qr:                    # entirely inside → use stored sum
        if probe is not None:
            probe.append({"node": node, "lo": lo, "hi": hi, "kind": "pick"})
        return self._tree[node]
    if probe is not None:                        # partial overlap → descend
        probe.append({"node": node, "lo": lo, "hi": hi, "kind": "split"})
    mid = (lo + hi) // 2
    return (self._query(2 * node, lo, mid, ql, qr, probe)
            + self._query(2 * node + 1, mid + 1, hi, ql, qr, probe))

Y el update repara el árbol recalculando los ancestros de un solo camino:

def update(self, i, value):
    """Set data[i] = value in O(log n): change the leaf, then recompute the stored
    aggregate of every ancestor on the path back to the root — there are only log n
    of them, so no full rebuild."""
    self._update(1, 0, self._n - 1, i, value)

def _update(self, node, lo, hi, i, value):
    if lo == hi:
        self._tree[node] = value
        self._data[i] = value
        return
    mid = (lo + hi) // 2
    if i <= mid:
        self._update(2 * node, lo, mid, i, value)
    else:
        self._update(2 * node + 1, mid + 1, hi, i, value)
    self._tree[node] = self._tree[2 * node] + self._tree[2 * node + 1]

Míralo funcionar

Aquí hay un segment tree sobre ocho elementos y una consulta por la suma del rango [2:6]. Cada nodo está etiquetado con el rango que cubre. Avanza paso a paso por la consulta: arranca en la raíz (que cubre 0:7, se traslapa parcialmente, así que se parte — naranja) y va bajando. Donde el rango de un nodo cae completamente dentro de [2:6], la consulta toma su suma guardada sin seguir bajando (verde — un segmento canónico); donde un nodo queda en parte dentro y en parte fuera, se vuelve a partir. Fíjate en qué tan pocos nodos verdes se necesitan para cubrir todo el rango: apenas tres segmentos canónicos arman la respuesta de 21, en lugar de visitar los cinco elementos. Ese puñado de piezas precalculadas es el O(log n):

El código completo

Las dos versiones en un solo lugar — cambia entre ellas. La pestaña desde cero es el segment tree. La pestaña de librería son las dos líneas base con array a las que les gana — el array en crudo (update rápido, consulta lenta) y las prefix sums (consulta rápida, update lento) — porque no hay segment tree en la librería estándar; es una estructura que construyes para el trabajo que tienes enfrente.

"""The segment tree — a tree over the positions of an array that answers RANGE queries
("sum, min, or max of everything between i and j") and POINT updates ("set element i")
both in O(log n).

A plain array forces a choice: keep the raw values and range queries cost O(n); keep
prefix sums and queries are O(1) but any update costs O(n) to rebuild. The segment tree
refuses the choice. Each node stores the aggregate of a contiguous range — the root the
whole array, its children the two halves, and so on down to single elements — so a query
assembles its answer from O(log n) precomputed pieces, and an update touches only the
O(log n) nodes on one root-to-leaf path.
"""


class SegmentTree:
    def __init__(self, data):
        self._n = len(data)
        self._data = list(data)
        self._tree = [0] * (4 * max(1, self._n))     # safe upper bound on node count
        if self._n:
            self._build(1, 0, self._n - 1)

    def _build(self, node, lo, hi):
        if lo == hi:
            self._tree[node] = self._data[lo]
            return
        mid = (lo + hi) // 2
        self._build(2 * node, lo, mid)
        self._build(2 * node + 1, mid + 1, hi)
        self._tree[node] = self._tree[2 * node] + self._tree[2 * node + 1]

    # region: query
    def query(self, ql, qr, probe=None):
        """Sum of data[ql..qr] in O(log n). At each node: if its range is entirely
        OUTSIDE the query, it contributes 0; if entirely INSIDE, it contributes its
        stored aggregate directly — no need to descend; if it PARTIALLY overlaps, split
        into its two children. Only O(log n) nodes ever land 'entirely inside', and their
        ranges are the canonical pieces the answer is assembled from."""
        return self._query(1, 0, self._n - 1, ql, qr, probe)

    def _query(self, node, lo, hi, ql, qr, probe):
        if qr < lo or hi < ql:                       # entirely outside
            return 0
        if ql <= lo and hi <= qr:                    # entirely inside → use stored sum
            if probe is not None:
                probe.append({"node": node, "lo": lo, "hi": hi, "kind": "pick"})
            return self._tree[node]
        if probe is not None:                        # partial overlap → descend
            probe.append({"node": node, "lo": lo, "hi": hi, "kind": "split"})
        mid = (lo + hi) // 2
        return (self._query(2 * node, lo, mid, ql, qr, probe)
                + self._query(2 * node + 1, mid + 1, hi, ql, qr, probe))
    # endregion

    # region: update
    def update(self, i, value):
        """Set data[i] = value in O(log n): change the leaf, then recompute the stored
        aggregate of every ancestor on the path back to the root — there are only log n
        of them, so no full rebuild."""
        self._update(1, 0, self._n - 1, i, value)

    def _update(self, node, lo, hi, i, value):
        if lo == hi:
            self._tree[node] = value
            self._data[i] = value
            return
        mid = (lo + hi) // 2
        if i <= mid:
            self._update(2 * node, lo, mid, i, value)
        else:
            self._update(2 * node + 1, mid + 1, hi, i, value)
        self._tree[node] = self._tree[2 * node] + self._tree[2 * node + 1]
    # endregion

"""There's no segment tree in the standard library — it's a build-it-yourself structure.
The instructive counterparts are the two array approaches it improves on, each fast at one
operation and slow at the other:

  - a RAW array: O(1) update, but O(n) range query (you sum the slice every time);
  - a PREFIX-SUM array: O(1) range query, but O(n) update (any change rebuilds the sums).

The segment tree's pitch is O(log n) for BOTH, so it wins on any workload that mixes
queries and updates. The face-off shows exactly that.
"""


# region: naive_array
class NaiveArray:
    """Raw values: updating is trivial, but every range query re-sums the slice — O(n)."""

    def __init__(self, data):
        self._d = list(data)

    def query(self, lo, hi):
        return sum(self._d[lo:hi + 1])          # O(n)

    def update(self, i, value):
        self._d[i] = value                      # O(1)
# endregion


# region: prefix_sums
class PrefixSums:
    """Precomputed prefix sums: a range query is one subtraction — O(1) — but any update
    invalidates the sums and forces an O(n) rebuild."""

    def __init__(self, data):
        self._d = list(data)
        self._rebuild()

    def _rebuild(self):
        self._p = [0]
        for x in self._d:
            self._p.append(self._p[-1] + x)

    def query(self, lo, hi):
        return self._p[hi + 1] - self._p[lo]    # O(1)

    def update(self, i, value):
        self._d[i] = value
        self._rebuild()                         # O(n)
# endregion

Desde cero vs librería

Este face-off es la mejor ilustración de todo el nivel de lo que significa emparejar una estructura con una carga de trabajo y no con una sola operación. En puras consultas, las prefix sums aplastarían al segment tree; en puros updates, lo haría el array en crudo. Pero mezcla las dos — como lo hace cualquier aplicación real — y cada enfoque con array queda arrastrado por la operación que dejó en O(n), mientras el segment tree, O(log n) en ambas, pasa de largo: 6× más rápido que el array en crudo, 100× más rápido que las prefix sums, en la misma mezcla de 5000 operaciones. La lección se generaliza mucho más allá de los segment trees: cuando hagas benchmark, haz benchmark de la carga de trabajo, porque una estructura óptima para una operación aislada puede ser pésima para la mezcla que realmente vas a correr. El segment tree gana no por ser el más rápido en algo, sino por no tener ninguna operación lenta.

Dónde te lo vas a encontrar de verdad

Los segment trees corren consultas de agregados por rango sobre datos que cambian. Los sistemas de analítica y monitoreo los usan (a ellos y a sus variantes) para estadísticas de rango móviles sobre series de tiempo mutables. Las bases de datos usan interval trees y range trees — parientes cercanos — para responder "qué filas caen en este rango" e indexar datos geométricos. La programación competitiva se apoya en ellos para todo un género de problemas de "consulta de rango con updates". La geometría computacional, su lugar de nacimiento, los usa para consultas de ventana y de stabbing. Y cualquier sistema que tenga que responder "el máximo/suma/mínimo sobre esta ventana deslizante o arbitraria, mientras los datos de abajo se actualizan" es candidato. Cuando conviven rangos y updates, el segment tree es la respuesta general.

Puntos clave

Un segment tree guarda el agregado de cada rango contiguo en un árbol binario, así que una consulta de rango arma su respuesta con $O(\log n)$ segmentos canónicos y un update puntual repara $O(\log n)$ ancestros — haciendo ambas rápidas donde un array plano tiene que sacrificar una. Funciona para cualquier operación asociativa, se extiende con lazy propagation a updates de rango, y es la herramienta correcta siempre que los datos mutables se topen con consultas de rango, a cambio de más código y más memoria que estructuras más simples.

El siguiente capítulo es justo esa estructura más simple, para el caso especial más común. Si tu agregado es una suma (o cualquier cosa invertible) y necesitas consultas de prefijo con updates puntuales, el Fenwick tree — también llamado binary indexed tree — hace el mismo trabajo en O(log n) con un puñado de líneas y un solo array, usando un truco precioso con la representación binaria de los índices. Es el hermano flaco y especializado del segment tree.