D - 気温の統一 / Uniform Temperature Editorial by admin
GPT 5.4 High概要
この問題は、総コスト
\(\sum_{i=1}^{N} S_i \times |P_i - T|\)
を最小にする整数 \(T\) を求める問題です。
重要なポイントは、これは「重み付き中央値(weighted median)」を使うと \(O(N \log N)\) で解けることです。
考察
各部屋 \(i\) について、目標温度を \(T\) にしたときのコストは
\(S_i \times |P_i - T|\)
です。したがって全体のコストは
\(f(T)=\sum_{i=1}^{N} S_i \times |P_i - T|\)
となります。
重要な気づき
この形は「絶対値の和」なので、最小値は中央値の考え方で求まります。
ただし今回は各部屋ごとに重み \(S_i\) があるので、普通の中央値ではなく 重み付き中央値 になります。
なぜ中央値になるのか
温度を小さい順に並べて考えます。
ある整数 \(T\) から \(T+1\) に目標温度を上げたとき、総コストがどれだけ変わるかを考えます。
- \(P_i \le T\) の部屋は、目標温度が 1 上がると距離が 1 増えるので、コストは \(S_i\) 増える
- \(P_i > T\) の部屋は、目標温度が 1 上がると距離が 1 減るので、コストは \(S_i\) 減る
よって、総コストの変化量は
\(\sum_{P_i \le T} S_i - \sum_{P_i > T} S_i\)
です。
つまり、
- 右側の重みのほうが大きいなら、\(T\) を右へ動かすとコストが下がる
- 左側の重みのほうが大きいなら、\(T\) を右へ動かすとコストが上がる
したがって最小になるのは、
- 左側の重みが全体の半分以下
- 右側の重みも全体の半分以下
となる位置です。これが 重み付き中央値 です。
素朴な方法がなぜだめか
目標温度 \(T\) は任意の整数なので、\(-10^9\) から \(10^9\) まで全部試すのは不可能です。
また、候補を各 \(P_i\) だけに絞ったとしても、毎回全体コストを計算すると \(O(N^2)\) になり、\(N \le 150000\) では間に合いません。
どう解決するか
温度 \(P_i\) でソートし、重み \(S_i\) の累積和を見ます。
全重みを \(W=\sum S_i\) とすると、
小さい温度から順に見ていって、累積重みが初めて \(W/2\) 以上になる温度が重み付き中央値です。
コードでは小数を使わず、
\(2 \times \text{acc} \ge W\)
を満たした最初の温度を選んでいます。
例えば、
- \((P,S)=(18,2),(20,5),(25,1)\)
なら全重みは \(8\) です。
温度順に累積すると
- \(18\) までで \(2\)
- \(20\) までで \(7\)
ここで初めて半分の \(4\) を超えるので、重み付き中央値は \(20\) です。
アルゴリズム
- 各部屋を \((P_i, S_i)\) として配列に入れる
- 温度 \(P_i\) で昇順ソートする
- 全重み \(W=\sum S_i\) を求める
- 小さい温度から順に累積重み
accを足していき、初めて \(2 \times acc \ge W\) となる温度をmedianとする - 最後に \(\sum S_i \times |P_i - \text{median}|\) を計算して出力する
この方法で、最小コストを求められます。
計算量
- 時間計算量: \(O(N \log N)\)
- 空間計算量: \(O(N)\)
実装のポイント
ソートしやすいように
(温度, 広さ)の順で保存している重み付き中央値の判定は
acc * 2 >= total_weightとして、浮動小数点を使わずに書いている答えは非常に大きくなる可能性があるので、Python 以外では 64 bit 整数型が必要になる場合がある
ソースコード
import sys
def main():
input = sys.stdin.readline
N = int(input())
rooms = []
total_weight = 0
for _ in range(N):
s, p = map(int, input().split())
rooms.append((p, s))
total_weight += s
rooms.sort()
acc = 0
median = 0
for p, s in rooms:
acc += s
if acc * 2 >= total_weight:
median = p
break
ans = 0
for p, s in rooms:
ans += s * abs(p - median)
print(ans)
if __name__ == "__main__":
main()
この解説は gpt-5.4-high によって生成されました。
posted:
last update: