AtCoder Beginner Contest 023 C: 収集王（片側全探索）

問題

abc023.contest.atcoder.jp

解法

kmjp.hatenablog.jp

コード

import java.util.Scanner;
import java.util.TreeSet;

public class Main {
    private static final int MAX = 100000;

    public static void main(String[] args) {
	Scanner sc = new Scanner(System.in);
	int R = sc.nextInt();
	int C = sc.nextInt();
	int K = sc.nextInt();
	int N = sc.nextInt();

	int[] col = new int[C];
	int[] row = new int[R];
	TreeSet<Long> candies = new TreeSet<>();
	for (int i = 0; i < N; i++) {
	    int r = sc.nextInt() - 1;
	    int c = sc.nextInt() - 1;

	    col[c]++;
	    row[r]++;

	    candies.add((long) r * MAX + c);
	}
	sc.close();

	int[] countCol = new int[MAX + 1];
	for (int i = 0; i < col.length; i++) {
	    countCol[col[i]]++;
	}

	long ans = 0;
	for (int i = 0; i < row.length; i++) {
	    if (K >= row[i]) {
		ans += countCol[K - row[i]];
	    }
	}

	for (Long long1 : candies) {
	    int r = (int) (long1 / MAX);
	    int c = (int) (long1 % MAX);

	    if (col[c] + row[r] == K) {
		ans--;
	    } else if (col[c] + row[r] == K + 1) {
		ans++;
	    }
	}
	System.out.println(ans);

    }

}