C - 収集王 〜ABC023〜

問題

C - 収集王

感想

飴があるところとないところを分けて考えることが難しかった。この分け方を思いつくまでにかなりの時間を要しました。

難しかったけれどその分解法を思いついた時の嬉しさはひとしおでしたね。

考え

まず，(i, j) について二重ループを回すことを思いついた。つまり，

• 各マス目を全探索する
• (i, j)のとき獲得する飴の個数は，各 i行，j 列についての累積和（総数）を先に計算しておいて求める。

ということだが，条件より各マスを探索するとTLEになってしまう。なので全探索をしないで求める必要がある。

次に思いついたのは，各マスを考えるのではなく，行と列に分けて求めるということである。つまり，

• i 行目，j 列目の飴の総数（sum_r[i], sum_c[j] とした）を記録しておく。
• 記録した後で，行と列ごとに飴の総数を記録する。（行についてみたときに飴の数が1個とれる行は sc[1]つ ，2個とれる行は sc[2]つ...，列についてみたときに飴の数が1個とれる列は sr[1]つ...というように）
• K個獲得するということは，行についてみたときに num 個とれて，列についてみたときに K - num 個とれるということに対応するので，K 回ループさせる。

ただしかし，ここで問題となるのは，飴がある場所については，sum_r[i]，sum_c[j] が二重でカウントされていることである。

なので，上記の手順で計算して求められた値 ans に対して，「飴がある場所に移動してかつ，総数が K - 1 であったマスの数」を引いて，「飴がある場所に移動してかつ，総数が `K であったマスの数」を足す必要がある。

Accepted Program

#include <map>
#include <iostream>
using namespace std;
#define ll long long
int main () {
  ll R, C , K;
  cin >> R >> C >> K;
  ll N;
  cin >> N;
  ll r[N], c[N];
  ll sum_r[R], sum_c[C];
  ll sr[R + 1], sc[C + 1];
  for (ll i = 0; i < R; i++) {
    sum_r[i] = 0;  
    sr[i] = 0;  
  }
  sr[R] = 0;
  for (ll i = 0; i < C; i++) {
    sum_c[i] = 0;  
    sc[i] = 0;  
  }
  sc[C] = 0;
  for (ll i = 0; i < N; i++) {
    cin >> r[i] >> c[i];
    r[i]--;
    sum_r[r[i]]++;
    c[i]--;  
    sum_c[c[i]]++;
  }
  for (ll i = 0; i < R; i++) {
    sc[sum_r[i]]++;
  }
  for (ll i = 0; i < C; i++) {
    sr[sum_c[i]]++;
  }
  ll ans = 0;
  for (ll i = 0; i <= K; i++) {
    if (K - i <= C && i <= R) {
    ans += sr[i] * sc[K - i];    
    }
  }
  ll plus, minus;
  plus = 0;
  minus = 0;
  for (ll i = 0; i < N; i++) {
    if (sum_r[r[i]] + sum_c[c[i]] - 1 == K - 1) {
      minus++;  
    }     
    if (sum_r[r[i]] + sum_c[c[i]] - 1 == K) {
      plus++;    
    }
  }
  ans = ans - minus + plus;
  cout << ans << endl;
}