burrows_wheeler_transform

burrows_wheeler_transform.py

def burrows_wheeler_transform(s: str) -> str:
    assert s.count("$") == 1
    n = len(s)
    table = [[None] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            table[i][j] = s[(i + j) % n]

    table.sort()
    return "".join([table[i][-1] for i in range(len(table))])


# L[i]が対応するF[x]のxを返す
def LF(i, L):
    # C[L[i]]
    a = list(sorted(L)).index(L[i])

    # occ(L[i], 0, i)
    b = sum([L[x] == L[i] for x in range(i)])
    return a + b


def inverse_burrows_wheeler_transform(L: str) -> str:
    assert L.count("$") == 1
    n = len(L)
    now = L.index("$")
    ans = ""
    for _ in range(n):
        ans += L[now]
        now = LF(now, L)

    return ans[::-1]


def main():
    s = "abraca$"
    print("inp:", s)
    bwt = burrows_wheeler_transform(s)
    print("bwt:", bwt)
    rev = inverse_burrows_wheeler_transform(bwt)
    print("inv:", rev)


if __name__ == '__main__':
    main()