ABC023 C- 収集王

問題

C - 収集王

検討

10^6*10^6と非常に大きな領域だけどデータセットは10^3以下なので必要なところだけ計算することを考えましたが，まだTLEします。

また，飴のあるマスは縦横に二重にカウントしているので，その部分を調整することを考えました。

しかしWAがとれません。
更に考えてみます。

WA/TLEしたコード

#include <bits/stdc++.h>
using namespace std;
int rnum[100100];
int cnum[100100];

int main(){
  int r,c,k;
  cin>>r>>c>>k;
  int n;
  cin>>n;
  vector<pair<int,int>>P;
  int a,b;
  for(int i=0;i<n;i++){
    cin>>a>>b;
    a--;
    b--;
    P.push_back({a,b});
    rnum[a]++;
    cnum[b]++;
  }
  int cnt=0;
  for(int i=0;i<r;i++){
    if(rnum[i] != 0){
      for(int j=0;j<c;j++){
        if(rnum[i]+cnum[j]==k){
          cnt++;
        }
      }
    }
  }
  for(int i=0;i<n;i++){ 
    if(rnum[P[i].first]+cnum[P[i].second]==k){
      cnt--;
    }
    if(rnum[P[i].first]+cnum[P[i].second]==k+1){
      cnt++;
    }
  }

  cout<<cnt<<endl;
  return 0;
}

更に検討

mapで各行と列の飴の個数を管理し，合計がkになるところの数を積で加算します。

飴のあるマスの調整は先の検討のとおりでOKでした。

int型だとオーバーフローしたのでlong long型にしました。

ACしたコード

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

#define int long long
int rnum[100100];
int cnum[100100];

signed main(){
  int r,c,k;
  cin>>r>>c>>k;
  int n;
  cin>>n;
  vector<pair<int,int>>P;
  int a,b;
  map<int,int>mpa;
  map<int,int>mpb;
  for(int i=0;i<n;i++){
    cin>>a>>b;
    a--;
    b--;
    P.push_back({a,b});
    rnum[a]++;
    cnum[b]++;
  }
  for(int i=0;i<r;i++){
    mpa[rnum[i]]++;
  }
  for(int i=0;i<c;i++){
    mpb[cnum[i]]++;
  }
  
  int cnt=0;
  for(int i=0;i<k+1;i++){
    cnt+=mpa[i]*mpb[k-i];
  }

  for(int i=0;i<n;i++){ 
    if(rnum[P[i].first]+cnum[P[i].second]==k){
      cnt--;
    }
    if(rnum[P[i].first]+cnum[P[i].second]==k+1){
      cnt++;
    }
  }

  cout<<cnt<<endl;
  return 0;
}