bintri.py

from collections import Counter
from itertools import product

def naive_encode(n):
    return '[' + '[]'*n + ']'

def encode(n):
    assert n == int(n) and n >= 0
    terms = []
    while n > 0:
        terms.append(encode(log2(n)))
        n -= 2**log2(n)
    return '[' + ''.join(terms) + ']'

def decode(xs):
    return sum(2**decode(i) for i in parse(xs))

def add(xs, ys):
    return xs[:-1] + ys[1:]

def mul(xs, ys):
    terms = (add(i, j) for i, j in product(parse(xs), parse(ys)))
    return '[' + ''.join(terms) + ']'

def canonicalize(xs):
    terms = map(canonicalize, parse(xs))
    iterate = True
    while iterate:
        new_terms = []
        iterate = False
        for i, j in Counter(terms).items():
            if j == 1:
                new_terms.append(i)
            else:
                fewer = (canonicalize(add(i, k)) for k in parse(encode(j)))
                new_terms.extend(fewer)
                iterate = True
        terms = new_terms
    return '[' + ''.join(sorted(terms)) + ']'

def log2(n):
    x = 0
    while 2**x <= n:
        x += 1
    return x - 1

def parse(xs):
    # parse("[[a][b][c]]") = ["[a]", "[b]", "[c]"]
    terms = []
    depth = 0
    for i, j in enumerate(xs):
        if j == '[':
            depth += 1
            if depth == 2:
                start = i
        if j == ']':
            depth -= 1
            if depth == 1:
                end = i+1
                terms.append(xs[start:end])
        assert depth > 0 or i == len(xs)-1, 'Bad string {}'.format(xs)
    assert depth == 0, 'Bad string {}'.format(xs)
    return terms

def run_tests():
    for i in range(1000):
        # No information is lost
        assert i == decode(encode(i))

        # Canonicalization is not needed after encode
        assert encode(i) == canonicalize(naive_encode(i))

        # Todo: randomize binary strings xs and assert that
        #       canonicalize(xs) == encode(decode(xs)).
    print('All OK')

def verify_youtube_examples():
    examples = (
        (1, ('[[]]',)),
        (2, ('[[][]]', '[[[]]]')),
        (3, ('[[][][]]', '[[[]][]]')),
        (4, ('[[][][][]]', '[[[]][[]]]', '[[[[]]]]')),
        (5, ('[[][][][][]]', '[[[]][[]][]]', '[[[[]]][]]')),
        (65536, ('[[[[[[]]]]]]',)),
        (2**65536, ('[[[[[[[]]]]]]]',))
    )
    for i, j in examples:
        print(i if i != 2**65536 else '2^65536', end=' ')
        for k in j:
            assert i == decode(k)
            print('=', k, end=' ')
        assert j[-1] == canonicalize(j[-1])
        print('(canonical)')

def scroll(n, delay=0):
    import time
    for i in range(n):
        print(encode(i))
        time.sleep(delay)