蓝桥杯机器人塔——python

题目如下：

X 星球的机器人表演拉拉队有两种服装，A 和 B。

他们这次表演的是搭机器人塔。

类似：

A

B B

A B A

A A B B

B B B A B

A B A B B A

队内的组塔规则是：

A 只能站在 AA 或 BB 的肩上。

B 只能站在 AB 或 BA 的肩上。

你的任务是帮助拉拉队计算一下，在给定 A 与 B 的人数时，可以组成多少种花样的塔。

输入描述
输入一行两个整数 M,NM,N（0<M,N<5000<M,N<500）,分别表示 A、B 的人数，保证人数合理性。

输出描述
要求输出一个整数，表示可以产生的花样种数。

输入输出样例
示例
输入

1 2
输出

3

题目解析

这道题该怎么做那？我们实际上发现只需要查明白第一行就可以直接往上推了。就不用议论之这么多了。但是第一行应该有多少个呐？机器人总数 cnt = M+N，那么可以做多少层？可以这样思考，假设有 num 层，这就是长度为 num 步长为 1 的等差数列，那么塔内机器人总数 cnt = num*(num+1)/2 = M+N，所以 num = sqrt((M + N) * 2)。所以我们就定下来了，接下来就是来寻找第一行了，然后再查看是否可以向上扩展，A，B的数目是否足够使用。
在这里遍历所有第一行的选项，我们使用递归的方式，同时用1表示A，2表示B。
然后设置一个检查代码，从而行开始计算所有的A，B，个数如果到了顶层还是和m,n相同，说明是可行的。

code

import math

def check():# 判断是否可以向上拓展
    a = 0
    b = 0
    tmp = row
    while tmp >0:
        for i in range(1, tmp+1):
            if cnt[i] == 1:
                a +=1
            else:
                b += 1
        for i in range(2, tmp+1):
            if cnt[i-1]==cnt[i]:
                cnt[i-1] = 1
            else:
                cnt[i-1] = 2
        tmp -=1
    if a == m and b == n:
        return True
    else:
        return False

def dfs(k):# 遍历所有第一排
    global res
    if k > row:
        if check():
            res +=1
        return
    cnt[k] = 1
    dfs(k+1)
    cnt[k] = 2
    dfs(k+1)

if __name__ == "__main__":
    m