Bit matrix transpose algorithm in Python

Transposes an n by n bit matrix in lg n rounds of block swaps with masks and shifts, O(log n) word operations when the matrix fits in one word.

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

bit_matrix_transpose.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bit-matrix-transpose/. The authors are listed in prov/SYYJVH.jsonl.


def transpose(a: list[int]) -> list[int]:
    """Transpose an n x n bit matrix stored one row to a word, in place.

    Bit c of a[r] is the entry in row r and column c. n = len(a) must be a
    power of 2 and every a[r] must be below 2**n. Returns a, which now holds
    the transpose. Cost: lg n rounds of n/2 masked swaps.
    """
    n = len(a)
    if n == 0 or n & (n - 1):
        raise ValueError("n must be a power of 2")
    if any(not 0 <= v < (1 << n) for v in a):
        raise ValueError("each row must fit in n bits")
    j = n // 2
    m = (1 << j) - 1  # @line L1
    while j >= 1:  # @line L2
        for k in range(n):  # @line L3
            if k & j:
                continue
            t = ((a[k] >> j) ^ a[k + j]) & m  # @line L4
            a[k + j] ^= t  # @line L5
            a[k] ^= t << j  # @line L6
        j //= 2
        m ^= m << j  # @line L7
    return a  # @line L8


def transpose8x8(x: int) -> int:
    """Transpose an 8 x 8 bit matrix in one 64-bit word, entry (r, c) at bit 8r + c."""
    for s, m in ((7, 0x00AA00AA00AA00AA), (14, 0x0000CCCC0000CCCC), (28, 0x00000000F0F0F0F0)):
        t = (x ^ (x >> s)) & m
        x = x ^ t ^ (t << s)
    return x

test_bit_matrix_transpose.py 62 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bit-matrix-transpose/. The authors are listed in prov/SYYJVH.jsonl.

import itertools
import random

import pytest

from bit_matrix_transpose import transpose, transpose8x8


def naive(a: list[int]) -> list[int]:
    """Oracle: copy the bits one at a time."""
    n = len(a)
    return [sum(((a[r] >> c) & 1) << r for r in range(n)) for c in range(n)]


def naive8x8(x: int) -> int:
    return sum(((x >> (8 * r + c)) & 1) << (8 * c + r) for r in range(8) for c in range(8))


def test_example() -> None:
    assert transpose([0b0011, 0b0101, 0b1001, 0b0000]) == [7, 1, 2, 4]


@pytest.mark.parametrize("n", [1, 2, 4])
def test_exhaustive_small(n: int) -> None:
    for rows in itertools.product(range(1 << n), repeat=n):
        assert transpose(list(rows)) == naive(list(rows))


@pytest.mark.parametrize("n", [8, 16, 32, 64])
def test_random_against_oracle(n: int) -> None:
    rng = random.Random(1000 + n)
    for _ in range(60):
        a = [rng.getrandbits(n) for _ in range(n)]
        b = transpose(list(a))
        assert b == naive(a)
        assert transpose(list(b)) == a


def test_identity_is_fixed() -> None:
    for n in [1, 2, 4, 8, 16, 32]:
        eye = [1 << i for i in range(n)]
        assert transpose(list(eye)) == eye


def test_bad_input() -> None:
    with pytest.raises(ValueError):
        transpose([1, 2, 3])
    with pytest.raises(ValueError):
        transpose([4, 0])


def test_one_word() -> None:
    assert transpose8x8(0xFF) == 0x0101010101010101
    assert transpose8x8(0x8040201008040201) == 0x8040201008040201
    rng = random.Random(8)
    for _ in range(4000):
        x = rng.getrandbits(64)
        assert transpose8x8(x) == naive8x8(x)
        assert transpose8x8(transpose8x8(x)) == x

What the next level needs

the tests did not run in this build

Read the full page about Bit matrix transpose algorithm