优化完美交换计算：移除Python循环以解决TLE问题

首先，咱们先理清楚问题的数学本质，这样就能彻底摆脱耗时的循环啦！你的代码超时是因为当n很大时，遍历数组累加的操作会变得非常慢，而实际上我们可以用纯数学计算直接得到结果，完全不需要循环。

问题分析回顾 #

首先，序列总和 S = n*(n+1)//2，只有当S是偶数时才存在完美交换（否则直接返回0），此时目标和为 h = S//2。我们需要分两种情况计算完美交换的数量：

1. 情况1：原序列存在前缀和等于h
   也就是存在某个k，使得1+2+...+k = h。这种情况下，交换任意两个都在前k项里的元素，或者任意两个都在后n-k项里的元素，交换后仍然存在前缀和等于h（因为前k项和不变，或者后n-k项和不变，而总和是2h，所以后n-k项和等于h意味着前n-k项和也等于h）。这种情况的交换数量是组合数 C(k,2) + C(n-k,2)，也就是 k*(k-1)//2 + (n-k)*(n-k-1)//2。

2. 情况2：原序列不存在前缀和等于h
   此时我们找到最大的k，使得1+2+...+k < h，计算差值 d = h - (1+2+...+k)。完美交换的条件是交换一对元素(y, y+d)，其中y在前k项里，y+d在后n-k项里（这样交换后前k项和会变成h）。满足条件的数对数量就是 min(k, n-d)（因为y必须≤k，且y+d≤n）。

优化后的代码 #

import sys
import math

def count_perfect_swaps():
    input = sys.stdin.read().split()
    t = int(input[0])
    for i in range(1, t+1):
        n = int(input[i])
        s = n * (n + 1) // 2
        if s % 2 != 0:
            print(0)
            continue
        h = s // 2
        # 计算最大的k使得k*(k+1)//2 <= h
        sqrt_val = math.isqrt(8 * h + 1)
        k = (sqrt_val - 1) // 2
        if k * (k + 1) // 2 == h:
            # 情况1：存在前缀和等于h
            res = k * (k - 1) // 2 + (n - k) * (n - k - 1) // 2
        else:
            # 情况2：不存在前缀和等于h
            d = h - k * (k + 1) // 2
            res = min(k, n - d)
        print(res)

if __name__ == "__main__":
    count_perfect_swaps()

对原代码的改进点 #

1. 去掉不必要的数组：你原来创建的数组a完全没用，因为序列是1到n的连续整数，不需要实际生成数组。
2. 替换循环为数学计算：用平方根计算直接找到k，避免了从后往前累加的循环，时间复杂度从O(n)降到O(1)，完全解决超时问题。
3. 修正计算逻辑：原代码中计算s1的公式是错误的，现在用组合数正确计算情况1的交换数量。
4. 批量读取输入：用sys.stdin.read().split()一次性读取所有输入，比逐行读取更快，尤其适合测试用例多的情况。

这样修改后，不管n多大，代码都能快速运行啦！

内容的提问来源于stack exchange，提问作者Akash nitter