优化HackerEarth数组查询问题Python代码：降时间内存消耗

原来的计算式可以通过代数拆分将O(N*M)的计算简化为O(1)，核心是把双重求和的式子拆解成四个预计算总和的组合：

公式推导 #

目标值v的定义是：

v = Σ（A[i] * B[j] * ((i+1)+(j+1))），其中i遍历A的所有元素，j遍历B的所有元素

将式子展开并提取公共因子：

((i+1)+(j+1)) = (i+1) + (j+1)
v = Σ(A[i]*(i+1)*B[j]) + Σ(A[i]*B[j]*(j+1))
v = [Σ(A[i]*(i+1))] * [Σ(B[j])] + [Σ(A[i])] * [Σ(B[j]*(j+1))]

我们只需要维护四个总和变量：

• SA：数组A的元素总和
• SA_pos：A中每个元素乘以其1-based位置的总和（数组下标为0-based，对应位置为下标+1）
• SB：数组B的元素总和
• SB_pos：B中每个元素乘以其1-based位置的总和

这样每次计算v只需要执行：v = SA_pos * SB + SA * SB_pos，时间复杂度O(1)。

处理交换操作 #

每次交换时，仅需更新对应的总和变量即可，所有操作均为O(1)时间：

1. tp=1：交换A[i-1]与B[j-1]
   设A中目标位置元素为a，B中目标位置元素为b：

   • 更新SA：SA = SA - a + b
   • 更新SA_pos：SA_pos = SA_pos - a*i + b*i（i为题目给出的1-based位置）
   • 更新SB：SB = SB - b + a
   • 更新SB_pos：SB_pos = SB_pos - b*j + a*j
   • 交换数组元素：A[i-1], B[j-1] = B[j-1], A[i-1]

2. tp=2：交换A[i-1]与A[j-1]
   设A中两个位置元素为a和b，对应1-based位置为i和j：

   • SA总和不变，无需更新
   • 更新SA_pos：SA_pos = SA_pos - a*i - b*j + b*i + a*j
   • 交换数组元素：A[i-1], A[j-1] = A[j-1], A[i-1]

3. tp=3：交换B[i-1]与B[j-1]
   设B中两个位置元素为a和b，对应1-based位置为i和j：

   • SB总和不变，无需更新
   • 更新SB_pos：SB_pos = SB_pos - a*i - b*j + b*i + a*j
   • 交换数组元素：B[i-1], B[j-1] = B[j-1], B[i-1]

完整代码示例 #

import sys

def main():
    input = sys.stdin.read().split()
    ptr = 0
    T = int(input[ptr])
    ptr += 1
    for _ in range(T):
        N = int(input[ptr])
        M = int(input[ptr+1])
        ptr += 2
        A = list(map(int, input[ptr:ptr+N]))
        ptr += N
        B = list(map(int, input[ptr:ptr+M]))
        ptr += M
        Q = int(input[ptr])
        ptr += 1
        
        # 初始化总和变量
        SA = sum(A)
        SA_pos = sum(a * (idx + 1) for idx, a in enumerate(A))
        SB = sum(B)
        SB_pos = sum(b * (idx + 1) for idx, b in enumerate(B))
        
        res = []
        # 记录初始v值
        res.append(SA_pos * SB + SA * SB_pos)
        
        for __ in range(Q):
            tp = int(input[ptr])
            i = int(input[ptr+1])
            j = int(input[ptr+2])
            ptr += 3
            
            if tp == 1:
                x, y = i-1, j-1
                a, b = A[x], B[y]
                SA = SA - a + b
                SA_pos = SA_pos - a*i + b*i
                SB = SB - b + a
                SB_pos = SB_pos - b*j + a*j
                A[x], B[y] = b, a
            elif tp == 2:
                x, y = i-1, j-1
                a, b = A[x], A[y]
                SA_pos = SA_pos - a*i - b*j + b*i + a*j
                A[x], A[y] = b, a
            elif tp == 3:
                x, y = i-1, j-1
                a, b = B[x], B[y]
                SB_pos = SB_pos - a*i - b*j + b*i + a*j
                B[x], B[y] = b, a
            
            # 记录交换后的v值
            res.append(SA_pos * SB + SA * SB_pos)
        
        # 输出结果
        print('\n'.join(map(str, res)))
        
if __name__ == "__main__":
    main()

复杂度分析 #

• 预处理阶段：O(N+M)计算初始总和
• 每个查询：O(1)时间处理交换和计算v
• 整体时间复杂度：O(T*(N+M+Q))，相比原三重循环的O(T*(N+M+QNM))，效率提升几个数量级
• 内存复杂度：O(N+M)存储数组，加上四个变量，内存消耗极低

内容的提问来源于stack exchange，提问作者Vishwaajeet Bharti