Algorithm Search⌘K

Offline RMQ in linear time with union-find (Arpa's trick) in Python

A linear-time offline method for answering range minimum queries by processing the array with a stack and union-find, using near-constant amortized merges.

File
arpa_offline_rmq.py, 40 lines
Level
I0 not checked yet
Toolchain
Python 3.14, pytest 9
Test framework
pytest
Run the tests
python3 -m pytest
arpa_offline_rmq.pyRawGitHub
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/arpa-offline-rmq/. The authors are listed in prov/2G2RRV.jsonl.

from collections.abc import Sequence
from typing import Any


def arpa_offline_rmq(a: Sequence[Any], queries: Sequence[tuple[int, int]]) -> list[int]:
    """Return, for each query (L, R), the position of the leftmost minimum of a[L..R].

    The queries are answered offline in one sweep from left to right with a
    monotonic stack and a union-find forest with path compression.
    """
    n = len(a)
    p = list(range(n))  # @line L1
    by_right: list[list[int]] = [[] for _ in range(n)]  # @line L2
    for k, (lo, hi) in enumerate(queries):
        if not 0 <= lo <= hi < n:
            raise ValueError("query out of range")
        by_right[hi].append(k)
    ans = [0] * len(queries)
    stack: list[int] = []  # @line L3

    def find(x: int) -> int:
        r = x  # @line L11
        while p[r] != r:
            r = p[r]
        while p[x] != r:  # @line L12
            nx = p[x]
            p[x] = r
            x = nx
        return r  # @line L13

    for i in range(n):  # @line L4
        while stack and a[stack[-1]] > a[i]:  # @line L5
            p[stack.pop()] = i  # @line L6
        stack.append(i)  # @line L7
        for k in by_right[i]:  # @line L8
            ans[k] = find(queries[k][0])  # @line L9
    return ans  # @line L10
test_arpa_offline_rmq.py 51 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/arpa-offline-rmq/. The authors are listed in prov/2G2RRV.jsonl.

import random

import pytest

from arpa_offline_rmq import arpa_offline_rmq


def brute(a: list[int], lo: int, hi: int) -> int:
    best = lo
    for t in range(lo, hi + 1):
        if a[t] < a[best]:
            best = t
    return best


def test_example() -> None:
    a = [3, 1, 4, 1, 5, 9, 2, 6]
    qs = [(0, 2), (2, 5), (5, 5), (2, 6), (4, 7), (0, 7)]
    assert arpa_offline_rmq(a, qs) == [1, 3, 5, 3, 6, 1]


def test_edges() -> None:
    assert arpa_offline_rmq([], []) == []
    assert arpa_offline_rmq([7], [(0, 0)]) == [0]
    assert arpa_offline_rmq([2, 2, 2], [(0, 2), (1, 2)]) == [0, 1]
    with pytest.raises(ValueError):
        arpa_offline_rmq([1, 2], [(1, 0)])
    with pytest.raises(ValueError):
        arpa_offline_rmq([1, 2], [(0, 2)])


def test_long_decreasing_chain() -> None:
    n = 100000
    a = list(range(n, 0, -1))
    qs = [(0, n - 1), (0, 0), (5, n // 2)]
    assert arpa_offline_rmq(a, qs) == [n - 1, 0, n // 2]


def test_random_against_scan() -> None:
    rng = random.Random(7)
    for _ in range(600):
        n = rng.randrange(1, 30)
        a = [rng.randrange(6) for _ in range(n)]
        qs = []
        for _ in range(rng.randrange(1, 20)):
            lo = rng.randrange(n)
            qs.append((lo, rng.randrange(lo, n)))
        assert arpa_offline_rmq(a, qs) == [brute(a, lo, hi) for lo, hi in qs]

What the next level needs

  • the tests did not run in this build

Read the full page about Offline RMQ in linear time with union-find (Arpa's trick)