母羊生羊问题解决

python

问题是这样的: 1只母羊,在第2年和第4年生下小母羊,在第5年死去,问N年后有多少母羊?

想了很久,终于有了一个思路,不过目前效率不高

给出代码:

 1a = [1]
 2def sheep(N):
 3    N = int(N)
 4    if N == 1:
 5        return a
 6    for n,k in enumerate(a):
 7        a[n] += 1
 8    for n,k in enumerate(a):
 9        if k == 2:
10            a.append(1)
11        if k == 4:
12            a.append(1)
13    return sheep(N-1)
14
15while True:
16    b = sheep(input('请输入年份:\n'))
17    c = [i for i in b if i < 5]
18    print('母羊总数: %d'%len(c))
19    a = [1]