Official

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\) です。

アルゴリズム

  1. 各部屋を \((P_i, S_i)\) として配列に入れる
  2. 温度 \(P_i\) で昇順ソートする
  3. 全重み \(W=\sum S_i\) を求める
  4. 小さい温度から順に累積重み acc を足していき、初めて \(2 \times acc \ge W\) となる温度を median とする
  5. 最後に \(\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: