ABC116の「D - Various Sushi」

問題概要

AtCoder Beginner Contest 116 の、D - Various Sushi。 解説や提出されたコード見ても？？？だったけど、自分で書いてみて理解できた。

解き方

1. 種類ポイントのことは一旦忘れる

2. 単にK個分、得点順にでかいものだけを食べていく
   答えがこの種類数を下回ることはないのでここから上の種類数を調べる。

3. 選ばれてない種類の中で一番得点が高い寿司と、2個以上選ばれた種類の中で一番得点が低い寿司を交換する
   これを繰り返すことで種類数別の最大得点が得られるのであとは比較するだけ

コード

typedef long long ll;

typedef pair<int, int> P;

bool used[100005];
ll scores[100005];

int main(int argc, char *argv[]) {
  ll n,k;
  cin>>n>>k;
  // 全部の寿司を得点の高い順に
  priority_queue<P> all_pq;
  // タイプが重複した寿司のうち、得点の低いもの集合
  priority_queue<P, vector<P>, greater<P> > less_pq;
  for(ll i=0;i<n;i++){
    int t,d;cin>>t>>d;
    all_pq.push({d,t});
  }

  // 種類ポイントのことは一旦忘れる
  // 単にK個分、得点順にでかいものだけを食べていく
  ll sum = 0;
  int cnt=0;
  for(ll i=0;i<k;i++){
    auto v = all_pq.top();all_pq.pop();
    if(!used[v.second]){
      used[v.second] = 1;
      cnt++;
    }else{
      less_pq.push(v);
    }
    sum += v.first;
  }
  scores[cnt]=sum;

  // 選ばれてない種類の中で一番得点が高い寿司と、2個以上選ばれた種類の中で一番得点が低い寿司を交換する
  while(!all_pq.empty()&&!less_pq.empty()){
    auto v = all_pq.top();all_pq.pop();
    if(used[v.second])continue;
    auto v2 = less_pq.top();less_pq.pop();
    used[v.second]=true;
    sum -= v2.first;
    sum += v.first;
    cnt++;
    scores[cnt] = sum;
  }
  ll res=0;
  for(ll i=0;i<=n;i++){
    if(!scores[i])continue;
    res=max(res, scores[i] + i*i);
  }
  cout<<res<<endl;

  return 0;
}

プログラミングコンテストチャレンジブック [第2版]　?問題解決のアルゴリズム活用力とコーディングテクニックを鍛える?