けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 360 C - Move It (4Q, 灰色, 250 点)

集計処理 Greedy 最小コスト 最適化問題 箱とボール 数列 NoviSteps4Q AtCoder AtCoder250点 ABC-C 灰色diff バケット 最適化の考察：「≦ K であること」と「= K となる実例」を示す 最大値や最小値に着目する 考察：独立に考える

面白い問題！

問題概要

箱 1, 2, \dots, N とボール 1, 2, \dots, N がある。ボール i は箱 A_{i} に入っていて、その重さは W_{i} である。

これから、ボールの入っている箱を移すことで、どの箱にもちょうど 1 個ずつボールが入っている状態にしたい。ボールを移すコストは、そのボールの重さである。

実現するための最小コストを求めよ。

制約

  • 1 \le N \le 10^{5}

考えたこと

たとえば、 N = 6、 A = (2, 2, 5, 2, 5, 4)、 W = (9, 11, 5, 3, 7, 6) の場合を考えてみよう。このとき、

  • 箱 2 に、重さが (3, 9, 11) のボール (小さい順)
  • 箱 4 に、重さが 6 のボール
  • 箱 5 に、重さが (5, 7) のボール (小さい順)

という状況になっている。まず、箱 4 については何もする必要がない。箱 2 については、3 個のボールのうち、重さが小さい 2 個のボールを他に移せばよい (移す先の箱はなんでもよい)。箱 5 についても、2 個のボールのうち、重さが小さい 5 のボールを他に移せばよい。

一般に、ボールが複数個ある箱について、重さが最大のものを除外したボールの重さの総和を求め、その箱ごとに総和をとればよい。

計算量は O(N) となる。

コード

このコードでは、各箱のボールの重さの最大値を求めるのにソートを用いているので、計算量は O(N \log N) となっている。

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N;
    cin >> N;
    vector<long long> A(N), W(N);
    for (int i = 0; i < N; ++i) cin >> A[i], --A[i];
    for (int i = 0; i < N; ++i) cin >> W[i];
    
    vector<vector<long long>> tot(N);
    for (int i = 0; i < N; ++i) tot[A[i]].push_back(W[i]);
    for (int v = 0; v < N; ++v) sort(tot[v].begin(), tot[v].end());
    
    long long res = 0;
    for (int v = 0; v < N; ++v) {
        for (int i = 0; i + 1 < tot[v].size(); ++i) res += tot[v][i];
    }
    cout << res << endl;
}