SWEA 26011. 정수들의 합

정수 N,K가 주어질 때, 1 이상 N 이하의 자연수 a, b, c, d 중 a + b - c - d = K 를 만족하는 정수 쌍 (a, b, c, d) 의 개수를 세는 문제이다. 입력으로 첫 번째 줄에 정수 N, K가 주어진다. (1 ≤ N ≤ 100000, -2(N-1) ≤ K ≤ 2(N-1))

 

처음에는 단순하게 갯수를 a+b, c+d로 나누어 세어 O(N^2)  형태로 구하려 했으나 N이 10^6까지 될 수 있어 시간 초과가 발생하였다. 이를 해결하기 위해 합 빈도수 수식으로 O(1)로 계산하였다.

1 ≤ a, b ≤ N 에서 a+b = s인 쌍의 수:

count(s) = 
  s - 1        (2 ≤ s ≤ N+1)
  2N - s + 1   (N+1 < s ≤ 2N)

 

 

결과 코드는 다음과 같다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Solution {
	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int tc = Integer.parseInt(br.readLine());
		StringBuilder sb = new StringBuilder();
		for(int test_case=0; test_case<tc; test_case++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
			int n = Integer.parseInt(st.nextToken()), k = Integer.parseInt(st.nextToken());

	        long answer = 0;
	        for (long s = 2; s <= 2 * n; s++) {
	            long cdSum = s - k;
	            answer += countPairs(s, n) * countPairs(cdSum, n);
	        }
	        sb.append(answer).append("\n");
		}
		br.close();
		System.out.print(sb);
	}
	
	private static long countPairs(long s, long N) {
        if (s < 2 || s > 2 * N) return 0;
        if (s <= N + 1) return s - 1;
        else return 2L * N - s + 1;
    }
}