嵌套列表的问题变动了一下应该怎么做呢？

Question

输入一个嵌套列表，再输入层数，求该层的元素个数。如果求第二层次的元素个数，但是第二层次还有第三层次嵌套，则嵌套列表算一个元素。

Answer 1

def cp(li,n):
    t = li
    for i in range(n-1):
        g = []
        for v in t:
            if type(v) is list:
                g.extend(v)
        t = g

    return len(t)

lis = [
    [
        [1,0,1,0],
        [1,0,1,0],
        [1,0,1,0]
    ],
    [
        [1,0,1,0],
        [1,0,1,0],
        [1,0,1,0]
    ]
]

print(cp(lis,1))
print(cp(lis,2))
print(cp(lis,3))6