yokobuttonの不定期で競技プログラミングをするブログ

不定期で解けた競技プログラミングコンテストの問題を載せています。

キーエンスプログラミングコンテスト2021-Nov. (AtCoder Beginner Contest 227) D - Project Planning

問題の要約
 N個の部署があり,i(1<=i<=N)番目の部署にはAi人の社員がいる。異なる部署に同じ社員はいない。
 1つのプロジェクトをK個の部署から1人ずつ選出して作り、ちょうどK人から構成したい。
 プロジェクトは最大でいくつ作れるか?1人が複数のプロジェクトに参加できない。
制約
 1<=K<=N<=2*10^5
 1<=Ai<=10^12
入力
 N K
 A1 A2 ... AN
考え方
 1,N個の部署の番号を1つずらす。
  1<=i<=N => 0<=i<=N-1
 2,プロジェクトの最大数を問われているので、部署の人数Aiはソートしてよい。
 3,人数の少ない部署か多い部署のどちらからか人を選出すればよさそう。
  簡単な例で確かめると、多い部署からの方が多くプロジェクトを作れそう。
  ただし、最適解かはわからない。
 4,これで実装してみる。
#include<iostream>
#include<vector>
#include<algorithm>

using namespace std;

int main(){
    long long N, K;
      cin >> N >> K;
      vector<long long> A(N);
      for(long long i = 0; i < N; i++){
        cin >> A[i];
    }
      sort(A.begin(), A.end(), greater<long long>());
      long long count = 0;
      while(A[K-1] != 0){
          count += A[K-1];
          for(long long i = 0; i < K; i++){
            A[i] -= A[K-1];
        }
          sort(A.begin(), A.end(), greater<long long>());
    }
  
      cout << count << endl;
  
    return 0;
}
 5,結果、間違っていたようだし,TLEもあった。
 6,考え方を変えてみる。
  1個のプロジェクトを作るには1*K人必要。
  2個のプロジェクトを作るには2*K人必要。
  P個のプロジェクトを作るにはP*K人必要。
  全ての部署の人数を合計したものをAsumとする。
  Asum < P*KならばP個のプロジェクトを作れない。
  Asum >= P*KでもP個のプロジェクトを作れるかわからない。
 7,Xsum < P*KならばP個のプロジェクトを作れない。
  Xsum >= P*KならばP個のプロジェクトを作れる。
  そのような数をXsumとする。
 8,もう一度考えてみる。
  1個のプロジェクトを作るには1*K人必要。部署は別々。
  2個のプロジェクトを作るには2*K人必要。部署は別々。
  P個のプロジェクトを作るにはP*K人必要。部署は別々。
  全ての部署からP人以下を集めて、P*K人を集めればP個のプロジェクトができる。集めなければできない。
  これはXsumの要件を満たす。
 9,これはプロジェクトの数をPとすれば、Xsum+=Σmin(Ai,P)で求められる。
 10,素直に実装すれば計算量がダメそう。でもやってみる。
#include<iostream>
#include<vector>

using namespace std;

const long long inf = 10e12+1;

int main(){
    long long N, K;
      cin >> N >> K;
      vector<long long> A(N);
      for(long long i = 0; i < N; i++){
        cin >> A[i];
    }
      long long Pmax = 0;
      for(long long P = 1; P < inf; P++){
        long long Xsum = 0;
          for(long long i = 0; i < N; i++){
            Xsum += min(A[i], P);
        }
          if(Xsum < P*K){
            Pmax = P-1;
              break;
        }else if(Xsum >= P*K){
              continue;
        }
    }
  
      cout << Pmax << endl;
  
    return 0;
}
 11,結果TLEになる。予想通り。
 12,また考えてみる。
  最大P個のプロジェクトが作れるとすると、1個,2個,...,P個のプロジェクトが作れる。
  P+1個からは作れない。
  これは二分探索が使えそう
 13,二分探索を使った実装をする。

実際のプログラム
#include<iostream>
#include<vector>

using namespace std;

bool f(vector<long long> A, long long P, long long K){
      long long Xsum = 0;
      for(long long i = 0; i < A.size(); i++){
        Xsum += min(A[i], P);
    }
      if(Xsum < P * K){
          return false;
    }else if(Xsum >= P * K){
          return true;
    }
}

int main(){
    long long N, K;
      cin >> N >> K;
      vector<long long> A(N);
      for(long long i = 0; i < N; i++){
        cin >> A[i];
    }
      long long ok = 0;
      long long ng = 1e18/K;
      while(ng-ok > 1){
        long long mid = (ok + ng) / 2;
          if(f(A, mid, K)){
            ok = mid;
        }else{
            ng = mid;
        }
    }
      cout << ok << endl;
    return 0;
}