D - 身長差の総和 / Sum of Height Differences 解説 by admin
Claude 4.5 OpusOverview
This problem asks us to efficiently compute the sum of height differences (absolute values) over all pairs formed by choosing 2 people from \(N\) students.
Analysis
Naive Approach
The simplest method is to compute \(|A_i - A_j|\) for every pair \((i, j)\) (\(i < j\)) and sum them up.
total = 0
for i in range(N):
for j in range(i+1, N):
total += abs(A[i] - A[j])
However, the number of pairs is \(\frac{N(N-1)}{2}\), and when \(N = 2 \times 10^5\), this requires approximately \(2 \times 10^{10}\) operations. This will result in TLE (Time Limit Exceeded).
Key Insight: Leveraging Sorting
By sorting the array, we can remove the absolute values.
If the sorted array satisfies \(A_0 \leq A_1 \leq \ldots \leq A_{N-1}\), then for \(i < j\) we always have \(A_i \leq A_j\), so: $\(|A_i - A_j| = A_j - A_i\)$
Further Optimization: Considering Each Element’s Contribution
After sorting, we consider how much each element \(A_i\) contributes to the overall sum.
\(A_i\): - When forming pairs with the \(i\) elements before it (smaller), it is on the subtracted-from side (contributes as \(+A_i\) a total of \(i\) times) - When forming pairs with the \(N-1-i\) elements after it (larger), it is on the subtracting side (contributes as \(-A_i\) a total of \(N-1-i\) times)
However, let’s think about it in an even simpler way.
Understanding Through a Concrete Example
Sorting \(A = [3, 1, 4]\) gives \([1, 3, 4]\).
For \(i = 2\) (value \(4\)), the sum of differences with the preceding elements \([1, 3]\) is: $\((4 - 1) + (4 - 3) = 2 \times 4 - (1 + 3) = 8 - 4 = 4\)$
In other words, for the \(i\)-th element: $\(i \times A_i - (A_0 + A_1 + \ldots + A_{i-1})\)$
This can be computed efficiently using a prefix sum.
Algorithm
- Sort the array \(A\)
- Initialize the prefix sum
prefix_sumto \(0\) - For each element \(A_i\) (\(i = 0, 1, \ldots, N-1\)):
- Add the sum of differences with all elements before \(A_i\), which is \(i \times A_i - \text{prefix\_sum}\), to
total - Add \(A_i\) to
prefix_sum
- Add the sum of differences with all elements before \(A_i\), which is \(i \times A_i - \text{prefix\_sum}\), to
- Output
total
Computation on a Concrete Example
\(A = [3, 1, 4]\) → After sorting: \([1, 3, 4]\)
| \(i\) | \(A_i\) | \(i \times A_i - \text{prefix\_sum}\) | total | prefix_sum |
|---|---|---|---|---|
| 0 | 1 | \(0 \times 1 - 0 = 0\) | 0 | 1 |
| 1 | 3 | \(1 \times 3 - 1 = 2\) | 2 | 4 |
| 2 | 4 | \(2 \times 4 - 4 = 4\) | 6 | 8 |
Answer: \(6\) (Verification: \(|3-1| + |4-1| + |4-3| = 2 + 3 + 1 = 6\) ✓)
Complexity
- Time complexity: \(O(N \log N)\) (sorting is the bottleneck; the subsequent processing is \(O(N)\))
- Space complexity: \(O(N)\) (for storing the array)
Implementation Notes
Since the maximum value of \(A_i\) is \(10^9\) and the maximum value of \(N\) is \(2 \times 10^5\), the total sum can be on the order of \(10^{14}\). In Python, there is no need to worry about integer overflow, but in languages like C++, you need to use the
long longtype.Instead of maintaining a separate array for the prefix sum, we save memory by managing it with a single variable.
Source Code
def solve():
N = int(input())
A = list(map(int, input().split()))
# ソートする
A.sort()
# 累積和を計算
# ソート後、各要素A[i]について、それより小さい要素との差の総和を効率的に計算
# A[i]とA[j] (j < i)のペアについて、A[i] - A[j]の総和を求める
# = i * A[i] - (A[0] + A[1] + ... + A[i-1])
# = i * A[i] - prefix_sum[i]
total = 0
prefix_sum = 0
for i in range(N):
# A[i]より前の要素との差の総和
# A[i]はi個の要素より大きい(ソート済みなので)
total += i * A[i] - prefix_sum
prefix_sum += A[i]
print(total)
solve()
This editorial was generated by claude4.5opus.