import sys from functools import lru_cache MOD = 10**9 + 7 def is_terminal(R1, R2, B1, B2): a = R1 + R2 b = B1 + B2 if a == 0 or b == 0: return True # Check if all red cards are the same if R1 == 0 and R2 == 0: return False # impossible since a > 0 red_attack = 2 if R1 == 0 else (1 if R2 == 0 else None) if red_attack is None: return False # Check blue cards if B1 == 0 and B2 == 0: return False # impossible since b >0 blue_attack = 2 if B1 == 0 else (1 if B2 == 0 else None) if blue_attack is None: return False return red_attack == blue_attack def main(): n, m = map(int, sys.stdin.readline().split()) red = list(map(int, sys.stdin.readline().split())) blue = list(map(int, sys.stdin.readline().split())) red_ones = sum(1 for x in red if x == 1) red_twos = n - red_ones blue_ones = sum(1 for x in blue if x == 1) blue_twos = m - blue_ones @lru_cache(maxsize=None) def compute_E(R1, R2, B1, B2): a = R1 + R2 b = B1 + B2 if a == 0 or b == 0: return 0 if is_terminal(R1, R2, B1, B2): return 0 numerator_total = a * b sum_rest_part = 0 p_self = 0 for r in [1, 2]: if r == 1: if R1 == 0: continue prob_r_num = R1 else: if R2 == 0: continue prob_r_num = R2 for s in [1, 2]: if s == 1: if B1 == 0: continue prob_s_num = B1 else: if B2 == 0: continue prob_s_num = B2 numerator = prob_r_num * prob_s_num new_R1, new_R2, new_B1, new_B2 = R1, R2, B1, B2 if r > s: if s == 1: new_B1 -= 1 else: new_B2 -= 1 elif r < s: if r == 1: new_R1 -= 1 else: new_R2 -= 1 else: # r == s pass is_new_term = is_terminal(new_R1, new_R2, new_B1, new_B2) if is_new_term: contribution = 0 else: if (new_R1 == R1) and (new_R2 == R2) and (new_B1 == B1) and (new_B2 == B2): p_self += numerator else: e_val = compute_E(new_R1, new_R2, new_B1, new_B2) sum_rest_part += (numerator * e_val) denominator_total = a * b - p_self numerator_total = a * b + sum_rest_part inv_denom = pow(denominator_total, MOD-2, MOD) return (numerator_total * inv_denom) % MOD result = compute_E(red_ones, red_twos, blue_ones, blue_twos) print(result % MOD) if __name__ == '__main__': main()
采用QwQ-32B Q6_K GGUF用llama.cpp AI生成。
为了简化prompt,我对prompt做了些修改,删除了测试样例以及对分数模运算的解释。
完整的prompt:以下是一道算法题,请你用尽可能低的时间和空间复杂度,尽可能短的代码,在符合题目要求的情况下用Python3完成这道题,并给出解题思路,\n代表换行:“小红这天又和小紫发起了一场对战。对战的规则如下:两人各有一些怪兽卡。每回合两人随机的从自己当前存活的怪兽卡中抽取一张发起决斗,战斗力低的怪兽卡死亡。如果两张怪兽卡战斗力相同,则无事发生。游戏会进行到‘如果抽卡,则 100% 的概率无事发生’或者‘有一方卡牌被用完’时结束。请你计算小红和小紫游戏进行回合数的期望。\n输入描述:\n第一行输入两个正整数n,m,分别代表小红和小紫的怪兽卡数量。\n第二行输入n个正整数,代表小红的每张怪兽卡战斗力。\n第三行输入m个正整数,代表小紫的每张怪兽卡战斗力。\n1≤n,m≤50\n每张怪兽卡的战斗力只会是1或者2。\n输出描述:\n一个整数,代表最终回合数期望对10^9+7取模的值。可以证明,最终的答案一定是个有理数,你只需要输出其对10^9+7取模的结果。”