Official

B - 教室の割り当て / Classroom Assignment Editorial by kyopro_friends


各日ごとに独立に、各教室に何人の人が参加しようとしているかを求めます。

各教室に何人の人が参加しようとしているかを、配列(C++ なら vector, python なら list など)で管理する場合、通常の実装では 1 日あたり \(\Omega(N)\) 時間、全体で \(\Omega(NM)\) 時間となり TLE となります。
連想配列・辞書 (C++なら map, python なら dict など)を用いることで、希望者が \(0\) 人の教室の考慮が不要となり、1 日あたり \(O(K_i\log K_i)\) 時間、 \(K=\sum_i K_i\) として全体で \(O(K \log K)\) 時間で求めることができます。

実装例 (C++)

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

int main(){
  int n, m;
  cin >> n >> m;
  vector<int> c(m);
  for(int i=0; i<m; i++) cin >> c[i];
  int ans=0;
  for(int i=0; i<n; i++){
    map<int,int> d;
    int k;
    cin >> k;
    for(int j=0; j<k; j++){
      int p;
      cin >> p;
      d[p-1]++;
    }
    for(auto[k, v]: d){
      if(v <= c[k]){
        ans += v;
      }
    }
  }
  cout << ans << endl;
}

実装例 (Python)

from collections import defaultdict
N, M = map(int,input().split())
C = list(map(int,input().split()))
ans = 0
for _ in range(N):
  d = defaultdict(int)
  K = int(input())
  P = list(map(int,input().split()))
  for p in P:
    d[p-1] += 1
  for k, v in d.items():
    if v <= C[k]:
      ans += v

print(ans)

posted:
last update: