蓝桥杯真题——机器人塔

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,N（0<M,N≤50）,分别表示 A、B 的人数，保证人数合理性。
输出描述
要求输出一个整数，表示可以产生的花样种数。
输入输出样例
示例
输入

1 2

输出
3
运行限制
最大运行时间：10s
最大运行内存: 512M

算法思想：dfs

代码实现——python

# 这种慢速递归函数适合使用 lru_cache



from functools import lru_cache

a, b = (int(i) for i in input().split())
# 层数
n = (a + b)
for n in range(100000000):
    if 2 * (a + b) == (n * (n + 1)):
        if 2 * (a + b) == (n * (n + 1)):
            break


@lru_cache(None)
def dfs(n, x, a, b):
    if n == 0:  # 如果当前情况机器人A或者B数量大于它的最大值

        return a == b == 0  # 计算完后如果A,B机器人个数符合输入，返回true

    # next右移一位，就实现了错位，然后与now自身取异或，就得到了我们需要的相邻两位取异或
    # 但是最高位那一个异或的结果是多余的，这一步是消掉最高位的数

    next1 = (x ^ (x >> 1)) & ((1 << n) - 1)
    cnt = bin(next1)[2:].count('1')  # count()方法可快速求出二进制下1的个数
    a = a - n + cnt
    b = b - cnt
    return dfs(n - 1, next1, a, b)


c = 0
# n层的最下层是n个机器人，摆法有2的n次方种，每一种可以转化为数字c
for m in range(0, 1 << n): 
    if dfs(n, m, a, b):
        c += 1
print(c)