Make a binary tree based on binary numbers

make_tree_for_range.py

import sys
import itertools

def next_power_of_two(v):
    # From Stanford bit-twiddling hacks: https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
    v -= 1
    v |= v >> 1
    v |= v >> 2
    v |= v >> 4
    v |= v >> 8
    v |= v >> 16
    v |= v >> 32
    v += 1
    return v

def make_tree_for_range(begin, end):
    assert end >= begin
    assert isinstance(begin, int)
    assert isinstance(end, int)
    assert begin >= 0

    print(r'\tikz [binary tree layout, nodes={draw,circle}, font=\sffamily, semithick]')

    orig_mask = next_power_of_two(begin ^ end) - 1

    def print_tree(cur, mask):
        left = max(cur &~ mask, begin)
        right = min(cur | mask, end)
        assert cur|mask >= begin
        if left == right:
            print('node{%d} ' % left, end='')
        else:
            print('node{%d--%d} ' % (left, right), end='')
            assert left <= right
        if cur|(mask>>1) >= begin and mask > 0:
            print('child { ', end='')
            print_tree(cur, mask >> 1)
            print('} ', end='')
        else:
            print('child [missing] ', end='')
        if cur | ((mask + 1) >> 1) <= end and mask > 0:
            print('child {', end='')
            print_tree(cur | ((mask + 1) >> 1), mask >> 1)
            print('} ', end='')
        else:
            print('child [missing] ', end='')

    print('\\', end='')
    print_tree(begin&~orig_mask, orig_mask)
    print(';')


def main():
    make_tree_for_range(int(sys.argv[1]), int(sys.argv[2]))

if __name__ == '__main__':
    main()