Positive Pairs - MojaCoder

MojaCoder に問題を投稿したので、解法を書きます。
AtCoder緑レベルです。

問題

https://mojacoder.vercel.app/users/magurofly/problems/positive_pairs

問題文

$N$ 個の整数 $A_1, \ldots, A_N$ が与えられます。
$1 \le i \lt j \le N$ となるすべての $(i, j)$ に対して、 $\max(0, A_i A_j)$ の総和を $10^9 + 7$ で割ったあまりを出力してください。

制約

• $1 \le N \le 2 \times 10^5$
• $|A_i| \le 10^9$
• 入力はすべて整数

解説

解説

1. 正負に分ける

$\max(0, A_i A_j)$ は、 $A_i A_j$ が $0$ 未満、つまり負になるなら無視してもよい事を意味します。
$A_i A_j$ が負にならない条件は、 $A_i$ と $A_j$ の両方が非負、もしくは両方が負であることです。
$A$ を非負と負の２つの配列 $A^+$ と $A^-$ に分けます。これには ${\rm O}(N)$ かかります。
$A^+$ と $A^-$ のそれぞれで $A_i A_j$ の総和を計算します。

2. $A_i A_j$ の総和を計算する

$\sum_{i=0}^{N-1} \sum_{j=i+1}^{N} A_i A_j$ を変形します。
$A_i$ に $j$ は含まれないので、内側の $\sum$ の外に出すことができます。
すると、 $\sum_{i=0}^{N-1} A_i \sum_{j=i+1}^{N} A_j$ となります。
また、 $\sum_{j=i+1}^{N} A_j$ は累積和を取ることで、 ${\rm O}(1)$ で計算することができます。
これによって、 ${\rm O}(N)$ で解くことができました。

コード

MOD = 10**9 + 7
N = int(input())
A = map(int, input().split())

Ap, An = [], []
for x in A:
  if x >= 0: Ap.append(x)
  else: An.append(x)

ans = 0
for A in [Ap, An]:
  N = len(A)
  A_sum = [0]
  for i in range(N-1): A_sum.append((A_sum[-1] + A[N - i - 1]) % MOD)
  for i in range(N): ans = (ans + A[i] * A_sum[N - i - 1]) % MOD

print(ans)