from itertools import count, accumulate, islice
from sel_n_primers import sel_n_primers

# solució com a funció generadora
def n_senars_tetraedrics1(n):
    t = 1
    i = 1
    while n > 0:
        if t % 2 == 1:
            yield t
            n = n - 1
        i = i + 1
        t = t + (i * (i +1)) //2
        
# solució amb itertools
def n_senars_tetraedrics2(n):
    it1 = accumulate(accumulate(count()))  # nombres tetraèdrics
    it2 = filter(lambda i: i%2 == 1, it1)
    it3 = islice(it2, n)
    return it3

# solució que usa la funció de l'altre apartat
def n_senars_tetraedrics3(n):
    it = accumulate(accumulate(count()))  # nombres tetraèdrics
    return sel_n_primers(n, lambda t:t%2== 1, it)

# Tria la solució que vols provar
n_senars_tetraedrics = n_senars_tetraedrics2