AtCoder Beginner Contest 255 D問題 ±1 Operation 2

問題

atcoder.jp

提出解答

atcoder.jp

問題の概要

長さ $N$ の整数列 $A=\left(A_i \right)_{i=1}^N$ が与えられる. $q=1,2, \dots, Q$ に対して, 次の質問に答えよ.

の要素をにするために以下の操作は最低でも ( 回以上) 何回必要か?
- 以上以下の整数を選び, 以下のうちのどちらかを実行する.
  - $A_i$ に $1$ 加算する.
  - $A_i$ から $1$ 減算する.

制約

$1 \leq N,Q \leq 2 \times 10^5$
$0 \leq A_i, X_q \leq 10^9$

解法1 (答えの導出)

まず, $X_q$ が与えられた場合の答えを考える. このとき, 操作は項ごとに独立であることに着目する. このとき, できる操作は「 $1$ 加算」か「 $1$ 減算」なので, 各項ごとにしなければならない操作は

$A_i \leq X_q$ ならば $1$ 加算を $X_q-A_i$ 回する.
$A_i \geq X_q$ ならば, $1$ 減算を $A_i-X_q$ 回する.

ということになり, 最小回数は $\left|X-A_i \right|$ である.

これを各項について考えることにより, 求めるべき解答は

$\displaystyle \sum_{i=1}^N \left|X-A\_i \right|$

である.

解答2 (高速化)

ここで, 解答について, $A$ の並び替えについて不変であることがわかる. よって, $A$ は並び替えてもよいことがわかる. そこで, $A$ を昇順に並び替えるとする. 以降では $A$ は昇順であるとする.

このとき, $A_0=-\infty, A_{N+1}=+\infty$ とすることにより, 各 $X_q$ に対して, $A_k \leq X_q \lt A_{k+1}$ となる $0$ 以上 $N$ 以下の整数 $k$ が唯一存在する.

このとき, $A$ が昇順であることに注意して,

$\displaystyle A_{\rm sum}:=\sum_{i=1}^N A_i, \quad S_k=\sum_{i=1}^k A_i$

とおくと,

$\displaystyle \begin{align} \sum_{i=1}^N \left|X_q-A_i \right| &=\sum_{i=1}^k (X_q-A_i )+\sum_{i=k+1}^N (A_i-X_q) \\ &=kX_q-\sum_{i=1}^k A_i+\left(\sum_{i=1}^N A_i-\sum_{i=1}^k A_i \right)-(N-k)X_q \\ &=kX_q-S_k+(A_{\rm sum}-S_k)-(N-k)X_q \\ &=(2k-N)X_q+(A_{\rm sum}-2S_k)\\ \end{align}$

となる.

ここで, $k$ を求める方法について, $A$ が単調増加であることに注意すると, 二分探索で求めることができる.

よって, 計算量は前計算としてソートに $O(N \log N)$ 時間, $A_{\rm sum}, S_0, S_1, \dots, S_N$ を求めるためにまとめて $O(N)$ 時間かかり, 各質問にたいして二分探索を利用することにより $O(\log N)$ 時間で解答できる. 以上から, 全体では $O((N+Q) \log N)$ 時間となる.