今回も yukicoder の問題紹介・解説を行っていきたいと思います！
$3$ 問目に紹介する問題はこちら！
概要
問題 ： No.1012 荷物収集
レベル ： ★★★
実行時間制限 ： 2000 ms
問題文へのリンク↓↓
問題文
数直線上に $N$ 個の荷物が並んでいます、荷物 $i (1 \leq i \leq N)$ は位置 $x_i$ にあり、重さは $w_i$ です。
荷物を別の位置に運ぶためには、$|$(移動後の位置) - (移動前の位置)$|$ $\times$ (荷物の重さ) と等しいコストが必要です。
$Q$ 個の値 $X_1, X_2, ... , X_Q$ が与えられます。各 $i (1 \leq i \leq Q)$ に対して、すべての荷物を位置 $X_i$ に運ぶための必要最低限の合計コストをそれぞれ求めて 1 行ごとに出力してください。
制約
- 入力は全て整数
- $1 \leq N, Q \leq 10^5$
- $1 \leq x_i \leq 10^9$
- $x_i \neq x_j \quad (i \neq j)$
- $1 \leq w_i \leq 10^4$
- $1 \leq X_i \leq 10^9$
- $X_i \neq X_j \quad (i \neq j)$
出力
各 $i (1 \leq i \leq Q)$ に対して、必要最低限の合計コストをそれぞれ 1 行ごとに出力してください。
最後に改行してください。
サンプル
サンプル1
2 1
1 10
10 4
5
60
位置 $1$ にある荷物を位置 $5$ に運ぶのに必要なコストは $(5 - 1) \times 10 = 40$ で、位置 $10$ にある荷物を位置 $5$ に運ぶのに必要なコストは $(10 - 5) \times 4 = 20$ なので合計でコストが $60$ かかります。
サンプル2
3 3
4 3
9 4
10 10
3 9 10
97
25
22
解説
ここから解説を行っていきます。
自力で解いて見たい方はここで一旦止めてください！
絶対値を外す
とりあえず絶対値を見たら外すというのは数学の基本です。必ずすぐに発想として出せるようになりましょう。
では、 $i$ 番目のクエリ $X_i$ に対して、式変形を考えていきましょう。
ソートしても問題はないのであらかじめ $N$ 個の荷物は $x_i$ が昇順にソートしておきます。
# 入力 (x_i, w_i)
xw = [list(map(int, input().split())) for _ in range(n)]
# x_i でソートする
xw.sort(key=lambda x: x[0])
すると、二分探索を用いることで、 $X_i$ が $N$ 個の荷物のうちの何番目の位置に来るのかがわかります。
Python の場合は、 bisect があるので便利ですね。
C++ だと lower_boundみたいなやつでしょうか？間違ってたらすみません...
x_t \leqq X_i \leqq x_{t + 1} (1 \leq t \leq n - 1)
上記を満たす $t$ を二分探索で探すことができれば、後はひたすら式変形を行っていきます。
式変形
では式変形を行っていきましょう。
あるクエリ $X_i (1 \leq i \leq Q)$ に対して、
\begin{align}
\sum_{k = 1}^{N} |X_i - x_k| \times w_k &= \sum_{k = 1}^{t} (X_i - x_k) \times w_k + \sum_{k = t + 1}^{N} (x_k - X_i) \times w_k \\
&= \sum_{k = 1}^{t} X_i w_k - \sum_{k = 1}^{t} x_k w_k + \sum_{k = t + 1}^{N} x_k w_k - \sum_{k = t + 1}^{N} X_i w_k \\
&= 2\sum_{k = 1}^{t} X_i w_k - 2\sum_{k = 1}^{t} x_k w_k + \sum_{k = 1}^{N} x_k w_k - \sum_{k = 1}^{N} X_i w_k \\
&= 2 X_i \sum_{k = 1}^{t} w_k - 2\sum_{k = 1}^{t} x_k w_k + \sum_{k = 1}^{N} x_k w_k - X_i \sum_{k = 1}^{N} w_k \\
\end{align}
このようにして式変形を行うことができました。
$\sum$ は $1$ からでないと計算がかなり煩雑になってしまうので、無理やり全て $1$ から開始できるように添字を合わせています。
また、各クエリごとに与えられる $X_i$ は定数扱いができるので $\sum$ の外に出すことができます。
$\sum_{k = 1}^{N} w_k$ と $ \sum_{k = 1}^{N} x_k w_k$ は前処理で累積和を取っておくことで、各クエリごとに O(1) で求めることができます。
このようにして、この問題を解くことが出来ました。
ACコード
from bisect import bisect_left
n, q = map(int, input().split())
xw = [list(map(int, input().split())) for _ in range(n)]
x = list(map(int, input().split()))
xw.sort(key=lambda x: x[0])
sort_list = [xw[i][0] for i in range(n)] # 二分探索を行うためのリスト（ w_i は必要ないので、 x_i のみのリストを作成）
accum_xw = [0] # x_i * w_i の 累積和テーブル
accum_w = [0] # w_i の累積和テーブル
# 累積和を計算する
for i in range(n):
accum_xw.append(accum_xw[-1] + xw[i][0] * xw[i][1])
accum_w.append(accum_w[-1] + xw[i][1])
for i in range(q):
t = bisect_left(sort_list, x[i])
# 式変形で求めた４項を計算する
W = 2 * x[i] * accum_w[t] # 上記の式の第１項
X = -2 * accum_xw[t] # 上記の式の第2項
Y = accum_xw[-1] # 上記の式の第3項
Z = -x[i] * accum_w[-1] # 上記の式の第4項
ans = W + X + Y + Z
print(ans)
最後に
最後まで読んできただきありがとうございました！
数学問題は解いてて飽きませんね（笑）
リクエストなどがあればぜひ声かけて下さい！ありがとうございました！