AtCoder 🟡 ABC437D Sum of Differences

題目的難度顏色使用 Luogu 上的分級，由簡單到困難分別為 🔴🟠🟡🟢🔵🟣⚫。

🔗 ABC437D Sum of Differences

rating: 550

Problem Statement

題意簡述

給定兩個長度分別為 $N$ 與 $M$ 的正整數序列 $A = (A_1, \dots, A_N)$ 與 $B = (B_1, \dots, B_M)$。
請計算 $\displaystyle \sum_{i=1}^{N} \sum_{j=1}^{M} |A_i - B_j|$ 的值，並對 $998244353$ 取模。

Constraints

約束條件

• $1 \leq N, M \leq 3 \times 10^5$
• $1 \leq A_i, B_j < 998244353$

---

思路：排序後二分搜尋 + 前綴和優化

這道題的核心在於如何快速計算大量的「絕對值差」。

1. 暴力法的瓶頸

最直覺的做法是寫兩層迴圈，遍歷所有 $(i, j)$ 對並計算 $|A_i - B_j|$。
然而，由於 $N, M$ 皆可達 $3 \times 10^5$，總運算次數會達到 $9 \times 10^{10}$，這顯然會超時。我們需要一個接近 $\mathcal{O}(N \log N)$ 或 $\mathcal{O}(N+M)$ 的解法。

2. 拆解絕對值

絕對值 $|x - y|$ 的定義取決於 $x$ 和 $y$ 的大小關係：

$$|x - y| = \begin{cases} x - y & \text{if } x \ge y \\ y - x & \text{if } x < y \end{cases}$$

為了消去絕對值符號，我們必須知道 $A_i$ 和 $B_j$ 誰大誰小。

3. 固定一邊，對另一邊排序

我們可以固定序列 $A$ 中的每一個元素 $A_i$，然後一次性計算它對所有 $B_j$ 的貢獻。
為了快速區分哪些 $B_j$ 比 $A_i$ 小、哪些比 $A_i$ 大，我們可以先將序列 $B$ 進行排序。

假設 $B$ 已經由小到大排好序。對於當前的 $A_i$，我們可以利用 二分搜尋 (Binary Search) 找到一個分界點 idx，使得：

• $B_1, B_2, \dots, B_{idx}$ 皆小於等於 $A_i$
• $B_{idx+1}, \dots, B_M$ 皆大於 $A_i$

4. 推導公式

對於固定的 $A_i$，總和可以拆成兩部分：

1. 比 $A_i$ 小的部分 ($B_1 \dots B_{idx}$)：

   $$\sum_{j=1}^{idx} |A_i - B_j| = \sum_{j=1}^{idx} (A_i - B_j) = \text{idx} \times A_i - \sum_{j=1}^{idx} B_j$$

2. 比 $A_i$ 大的部分 ($B_{idx+1} \dots B_M$)：

   $$\sum_{j=idx+1}^{M} |A_i - B_j| = \sum_{j=idx+1}^{M} (B_j - A_i) = \sum_{j=idx+1}^{M} B_j - (M - \text{idx}) \times A_i$$

5. 前綴和加速計算

上述公式中，$\sum B_j$ 的部分如果是每次重新加總，複雜度依然太高。
我們可以預先計算 $B$ 的 前綴和 (Prefix Sum) 陣列 sb，其中 sb[k] 代表 $\sum_{j=1}^k B_j$。
如此一來：

• $\sum_{j=1}^{idx} B_j$ 就是 sb[idx]。
• $\sum_{j=idx+1}^{M} B_j$ 就是 sb[M] - sb[idx]。

這兩個操作都可以在 $\mathcal{O}(1)$ 時間內完成。

模運算注意

題目要求對 $998244353$ 取模。雖然 Python 自動處理大整數，但在計算過程中隨時取模可以保持數值較小。
另外，利用前綴和相減計算區間和時，若在其他語言中可能會出現負數，需注意 (a - b + MOD) % MOD 的處理。

複雜度分析

• 時間複雜度：$\mathcal{O}((N+M) \log M)$
  • 對 $B$ 排序：$\mathcal{O}(M \log M)$。
  • 計算 $B$ 的前綴和：$\mathcal{O}(M)$。
  • 遍歷 $A$ 並對每個元素進行二分搜尋：$\mathcal{O}(N \log M)$。
• 空間複雜度：$\mathcal{O}(N+M)$，用於儲存陣列與前綴和。

---

Code

from bisect import bisect_right
from itertools import accumulate

MOD = 998244353

def solve():
    n, m = map(int, input().split())
    A = list(map(int, input().split()))
    B = list(map(int, input().split()))
    assert len(A) == n and len(B) == m

    B.sort()
    sb = list(accumulate(B, initial=0))

    ans = 0
    for x in A:
        idx = bisect_right(B, x)
        ans += idx * x - sb[idx]
        ans += sb[m] - sb[idx] - (m - idx) * x
        ans %= MOD
    print(ans)

if __name__ == "__main__":
    solve()

---

寫在最後

PROMPT

這裡什麼都沒有。