subtree_sum.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Dados sequenciais
# dados = range(1, 1 << 5)

# Dados aleatórios duplicados
from random import randint
dados = [randint(0, 9) for x in xrange(1, 1 << 7)]

# Níveis
niveis = None


def mostrar_arvore(arvore):
    maior = max(reduce(lambda x, y: x + y, arvore))
    t = len(str(maior))
    ultima = ''
    for n in xrange(1, len(arvore) + 1):
        ultima = ' '.join(str(f).center(t) for f in arvore[-n]) + '\n' + ultima
        t = t * 2 + 1
    print ultima


# Preechimento de nível
def vetor_para_niveis(dados):
    global niveis
    niveis = [dados[0:1]]
    n = 1
    p = 1
    while p < len(dados):
        niveis.append(dados[p: p + (1 << n)])
        p += 1 << n
        n += 1


# Encontra o nível e índice de um dado valor
def encontrar_valor(valor, contagem=1):
    # Encontrar o nível e índice do valor dado
    c = contagem
    for n in range(len(niveis)):
        encontrados = niveis[n].count(valor)
        if encontrados:
            if encontrados < c:
                c -= encontrados
                continue
            i = -1
            while c:
                i = niveis[n].index(valor, i + 1)
                c -= 1
            return n, i
    raise ValueError("%iª ocorrência de %i inexistente!" % (contagem, valor))


# Somar o que estiver o próprio e abaixo de um nível n e de um índice i
# Intervalo de valores para cada nível inferior dado por i*2**p:(i+1)*2**p
def soma_sub_tab(n, i):
    soma = niveis[n][i]
    sub = [[soma]]
    tabela = ["   Acumulado Subtotal  Qtd", "%12i %8i %4i" % (soma, soma, 1)]
    for sn in range(n + 1, len(niveis)):
        p = sn - n
        valores = niveis[sn][i*2**p:(i+1)*2**p]
        s = sum(valores)
        soma += s
        sub.append(valores)
        tabela.append("%12i %8i %4i" % (soma, s, len(valores)))
    return soma, sub, tabela

# Transformar os dados em estrutura de níveis 2^n
print "Níveis"
vetor_para_niveis(dados)
mostrar_arvore(niveis)

# Pegar um valor aletório em dados e encontrar sua 2ª ocorrência
num = dados[randint(0, len(dados))]
n, i = encontrar_valor(num, 2)
print "2ª ocorrência de", num, "no nível", n + 1, "e índice", i + 1

# Somar a sub-árvore desse elemento a ele mesmo
print "\nIntervalos e soma da sub-árvore"
soma, sub_arvore, tabela = soma_sub_tab(n, i)
mostrar_arvore(sub_arvore)

print 'Tabela de cálculo'
print '\n'.join(tabela)

print "\nTotal:", soma