Zenn
⛳

ABC227 D - Project Planning C++解答例

2021/11/14に公開
C++
競技プログラミング
AtCoder
競プロ
tech

ABC227 D - Project PlanningをC++で解きます。

問題

問題文をAtCoderのページより引用します。

問題文

キーエンスにはN個の部署があり、i(1 \leq i \leq N)番目の部署にはA_i人の社員が所属しています。異なる部署に同じ社員が所属していることはありません。

キーエンスは、部署をまたいだ全社横断プロジェクトを計画しています。1つのプロジェクトはK個の相異なる部署から1人ずつ選出して作り、ちょうどK人から構成されるようにします。

プロジェクトは最大でいくつ作れますか？ただし、1人が複数のプロジェクトに参加することはできません。

制約

  • 1 \leq K \leq N \leq 2 \times 10^5
  • 1 \leq A_i \leq 10^{12}
  • 入力はすべて整数

解答例

二分探索による最大値探索を行う

制約条件を見ると、1つの部署に所属している社員の人数A_iの上限が10^{12}となっています。
プロジェクトを作るにはK個の異なる部署から1人ずつ社員を選んで作ることになりますが、この制約より、1人ずつ選ぶ処理を愚直に実装すると、最悪10^{12}回の計算を行わなければならなくなります。
(例えばN = 2, K = 2, A_1 = 10^{12}, A_2 = 10^{12}のような入力だったとすると、プロジェクトは10^{12}個作れる。)

この問題で求められている答えは、作成可能なプロジェクトの最大個数です。愚直に解くことはできない、かつ実現可能な最大値(最小値)を求めるような問題は、二分探索によって解くことを思いつきます。

作成するプロジェクトの個数を決め打ちする

求めるのはプロジェクトの最大個数なので、これを決め打ちする二分探索をすることを考えます。
作成するプロジェクトの個数をMと表すことにします。

二分探索の基本骨子は以下のようなコードになります。

二分探索の基本骨子

int64_t left = 0;
int64_t right = 適当な最大値;

while (right - left > 1) {
  int64_t mid = (left + right) / 2;
  
  bool ok = true;
  ~~何かしらの条件判定~~
  
  if (ok) {
    left = mid;
  } else {
    right = mid;
  }
}

ここでの何かしらの条件判定とは、「M個のプロジェクトを作成可能か？」です。この条件判定を実装すればこの問題を解くことができます。

M個のプロジェクトを作成できる条件

M個のプロジェクトを作成できる条件を考えます。
今、M個のプロジェクトがあり、各プロジェクトにはK人の社員がアサインされます。よって、以下のような図を作って考えてみます。

図のM個ある長方形が1つのプロジェクトを表しています。この長方形の中に、各部署に所属している社員を割り振っていくことを考えます。
ここで、ある部署iに所属しているA_i人の社員について考えてみます。同じプロジェクトに同じ部署の社員は所属することはできないので、各プロジェクトに1人ずつ置いていってみます。
もしA_iがM以下であれば、A_i人の社員を1人ずつプロジェクトにアサインすることができます。

一方で、もしA_iがMより大きければ、同じプロジェクトに同じ部署の社員は所属できないという制約より、A_i人いる社員のうち多くてもM人までしか、その部署から選ぶことができないことがわかります。

これを数式で表してみます。部署iからプロジェクトへ引き入れられる人数をB_iとすると、B_iは次の式で表されます。

B_i = \min(A_i, M)

N個ある各部署について、このB_i(1 \leq i \leq N)を計算すると、各部署からM個のプロジェクトへアサイン可能な最大人数が求められます。

M個のプロジェクトを作るとき、少なくとも全部でM \times K人の社員が必要となります。
つまり、各部署のB_iの総和が、M \times K未満ならば、M個のプロジェクトを作ることはできない、と判断することができます。
逆に\sum_{i=1}^N B_iがM \times K以上であれば、M個のプロジェクトを作成可能です(先の図の長方形全てを埋めることができるということ)。

二分探索の探索範囲

プロジェクトが作成可能かどうかを判定する条件ができたので、二分探索を実装します。
判定条件をプログラムで実装すると、次のようなコードになります。

二分探索の実装

int64_t left = 0;
int64_t right = 適当な最大値;

while (right - left > 1) {
  int64_t mid = (left + right) / 2;
  
  int64_t sum = 0;
  for (int64_t i = 0; i < N; i++) {
    sum += std::min(mid, A.at(i));
  }
  
  if (sum >= K * mid) {
    left = mid;
  } else {
    right = mid;
  }
}

二分探索の探索範囲(右側)ですが、Kを掛けたときにオーバーフローしない大きめの値を取っておきます。64bit符号付き整数の最大値はおよそ9 \times 10^{18}なので、9 \times 10^{18} \div Kくらいの値にしておきます。10^{18} \div Kでも大丈夫です。

実装例

コード全体は以下のようになります。

d.cpp
#include <iostream>
#include <vector>
#include <string>
#include <algorithm>

int main() {
  int64_t N, K;
  std::cin >> N >> K;
  std::vector<int64_t> A(N);
  for (int64_t i = 0; i < N; i++) {
    std::cin >> A.at(i);
  }
 
  int64_t left = 0;
  int64_t right = 9000000000000000000 / K;
 
  while (right - left > 1) {
    int64_t mid = (left + right) / 2;
 
    int64_t sum = 0;
    for (int64_t i = 0; i < N; i++) {
      sum += std::min(mid, A.at(i));
    }
 
    if (sum >= K * mid) {
      left = mid;
    } else {
      right = mid;
    }
  }
 
  std::cout << left << std::endl;
 
  return 0;
}

実際に提出したコードはこちら。

Discussion