Counting inversions with a Fenwick tree in Python

Scan left to right and count larger earlier values with a binary indexed tree.

File: inversion_count_with_fenwick_tree.py, 60 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

inversion_count_with_fenwick_tree.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/inversion-count-with-fenwick-tree/. The authors are listed in prov/XH0RVP.jsonl.

from bisect import bisect_left
from typing import Protocol, Sequence, TypeVar


class Ordered(Protocol):
    def __lt__(self, other: "Ordered", /) -> bool: ...


T = TypeVar("T", bound=Ordered)


def prefix(tree: list[int], i: int) -> int:
    """Return the number of recorded items with rank 1..i (tree is 1-based)."""
    s = 0  # @line L8
    while i > 0:  # @line L9
        s += tree[i]
        i -= i & -i
    return s  # @line L10


def add(tree: list[int], m: int, i: int, d: int) -> None:
    """Record d more items of rank i, for 1 <= i <= m."""
    while i <= m:  # @line L11
        tree[i] += d
        i += i & -i


def count_inversions_bit(a: Sequence[T]) -> int:
    """Return the number of pairs p < q with a[p] > a[q].

    Equal items are not counted. The values are compressed to ranks 1..m first.
    Time O(n log n), extra space O(n). The input is not changed.
    """
    v = sorted(set(a))  # @line L1
    m = len(v)
    r = [bisect_left(v, x) + 1 for x in a]  # @line L2
    tree = [0] * (m + 1)  # @line L3
    c = 0
    for j in range(1, len(a) + 1):  # @line L4
        c += (j - 1) - prefix(tree, r[j - 1])  # @line L5
        add(tree, m, r[j - 1], 1)  # @line L6
    return c  # @line L7


def prefix_inversion_counts(a: Sequence[T]) -> list[int]:
    """Return the list whose entry j is the inversion count of a[:j+1] (the online view)."""
    v = sorted(set(a))
    m = len(v)
    tree = [0] * (m + 1)
    c = 0
    out: list[int] = []
    for j, x in enumerate(a):
        k = bisect_left(v, x) + 1
        c += j - prefix(tree, k)
        add(tree, m, k, 1)
        out.append(c)
    return out

test_inversion_count_with_fenwick_tree.py 38 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/inversion-count-with-fenwick-tree/. The authors are listed in prov/XH0RVP.jsonl.

import random

from inversion_count_with_fenwick_tree import count_inversions_bit, prefix_inversion_counts


def brute(a: list) -> int:
    return sum(1 for i in range(len(a)) for j in range(i + 1, len(a)) if a[i] > a[j])


def test_example() -> None:
    assert count_inversions_bit([3, 1, 6, 2, 5, 4]) == 6
    assert prefix_inversion_counts([3, 1, 6, 2, 5, 4]) == [0, 1, 1, 3, 4, 6]


def test_repeated_values() -> None:
    assert count_inversions_bit([40, 10, 40, 30, 10]) == 6
    assert count_inversions_bit([7, 7, 7]) == 0
    assert count_inversions_bit([2, 2, 1]) == 2


def test_edges() -> None:
    assert count_inversions_bit([]) == 0
    assert count_inversions_bit([1]) == 0
    assert count_inversions_bit(list(range(10))) == 0
    assert count_inversions_bit(list(range(10, 0, -1))) == 45
    assert count_inversions_bit([1.5, -2.0, 10**18, -(10**18)]) == 4
    assert count_inversions_bit(["b", "a", "c"]) == 1


def test_random_against_brute_force() -> None:
    rng = random.Random(9103)
    for _ in range(600):
        a = [rng.randint(-5, 15) for _ in range(rng.randint(0, 40))]
        assert count_inversions_bit(a) == brute(a)
        assert prefix_inversion_counts(a) == [brute(a[: j + 1]) for j in range(len(a))]

What the next level needs

the tests did not run in this build

Read the full page about Counting inversions with a Fenwick tree