BOJ 백준 13553 수열과 쿼리 8

문제 : https://www.acmicpc.net/problem/13553

[ 문제 요약 ]

• 배열이 주어지면, 아래 조건에 해당하는 쌍의 개수를 출력합니다.
• l r: l ≤ i < j ≤ r이면서 abs(Ai - Aj) ≤ K인 (i, j) 쌍의 개수를 출력한다.

 

[ 테스트 케이스 설명 ]

4 31		// 수열크기 1<=100,000 , 정수K 1<=100,000
1 16 32 64	// 주어진 수 1<=100,000
4		// 쿼리 수 1<=100,000
1 4
1 2
2 4
2 3

 

 

[ 알고리즘 분류 ]

• 펜윅 트리, 오프라인 쿼리
• Mo’s 알고리즘, 자료구조

 

[ 문제 해설 ]

 

각 쿼리 당 구간이 주어지면, 해당 구간에 대해서 하나하나 처음부터 끝까지 탐색한다면 시간 초과를 받게 됩니다.

 

그래서 주어지는 질의 쿼리를 오프라인 쿼리로 처리하고, Mo’s 알고리즘으로 정렬한 후 투포인터 방식으로 풀어야 합니다.

 

각 구간의 쌍의 개수를 빠르게 가져오는 방법은 구간 합 세그먼트 트리를 이용하는 방법입니다.

 

구간 L, R이 주어지고 K 가주어진다면, K와 R을 이용해 L이 될 수 있는 범위를 구할 수 있습니다.

 

L에 대한 정리 식 : arr[R] - K ≤ arr[L] ≤ arr[R] + K

 

위와 같이 구간을 알고 특정수(arr[L])의 범위를 알 수 있기 때문에 구간 합 세그먼트 트리로 쌍을 구할 수 있습니다.

 

세그먼트 트리의 인덱스는 1부터 10만까지를 의미하고 그 안에 들어가는 값은, 해당 숫자가 나온 횟수(cnt)가 됩니다.

 

주어지는 수의 최댓값이 100,000이기 때문에 세그먼트 트리로 해도 메모리는 초과되지 않습니다.

 

그리고 쌍을 모두 더하는 것이기 때문에 int를 초과할 수 있어 결과는 long으로 선언해야 합니다.

 

그런데 이 문제는 자바로 푼다면, 세그먼트 트리로 구현 시 시간 초과를 받아 풀 수 없습니다.

 

세그먼트 트리는 재귀 오버헤드가 있기 때문에 이 문제에서는 오버헤드가 없는 펜윅 트리로 풀어야만 합니다.

 

[ 정답 코드 ]

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.StringTokenizer;
class Main{
	
	static int len = 100_000;
	static int fenwickTree[];
	
	public static void main(String[] args)throws Exception{
		BufferedReader	br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		int N		= Integer.parseInt(st.nextToken());// 수열크기 1<=100,000
		int K		= Integer.parseInt(st.nextToken());// 정수K 1<=100,000
		int arr[]	= new int[N + 1];
		fenwickTree		= new int[len + 1];
		
		st = new StringTokenizer(br.readLine());
		for(int i=1; i<=N; i++)
			arr[i] = Integer.parseInt(st.nextToken());
		
		ArrayList<Query> query = new ArrayList<>();
		int Q = Integer.parseInt(br.readLine());// 쿼리 수 1<=100,000
		long ans[] = new long[Q + 1];
		
		for(int i=1, sqrt = (int)Math.sqrt(N); i<=Q; i++)
		{
			st = new StringTokenizer(br.readLine());
			int l = Integer.parseInt(st.nextToken());
			int r = Integer.parseInt(st.nextToken());
			query.add(new Query(l, r, i, l / sqrt));
		}
		
		Collections.sort(query);
		
		int left = 1;
		int right = 0;
		long pairCnt = 0;
		for(Query q : query)
		{
			while(right < q.right)
			{
				++right;
				// right를 하나 넣으면 이전에 |v-x| <= K인 것들의 쌍이 한개씩 더생김
				// 그래서 arr[right]와 K를 활용해 나올 수 있는 범위의 개수를 모두 카운팅해서 더한다.
				pairCnt += sum(arr[right] - K, arr[right] + K);
				// 해당 값의 개수를 1추가한다.
				update(arr[right], 1);
			}
			
			while(q.right < right)
			{
				// 해당 값의 개수를 -1한다.
				update(arr[right], -1);
				// arr[right]가 하나 없어졌기 때문에, |v-x| <= K인 것들의 쌍이 하나씩 없어짐
				// 그래서 arr[right]와 K를 활용해 나올 수 있는 범위의 개수를 모두 카운팅해서 뺀다.
				pairCnt -= sum(arr[right] - K, arr[right] + K);
				--right;
			}
			
			while(left < q.left)
			{
				update(arr[left], -1);
				pairCnt -= sum(arr[left] - K, arr[left] + K);
				++left;
			}
			
			while(q.left < left)
			{
				--left;
				pairCnt += sum(arr[left] - K, arr[left] + K);
				update(arr[left], 1);
			}
			
			ans[q.idx] = pairCnt;
		}
		
		StringBuilder sb = new StringBuilder();
		
		for(int i=1; i<=Q; i++)
			sb.append(ans[i]).append('\n');
		
		System.out.print(sb);
	}
	static int sum(int l, int r) {
		if(l > r)
			return 0;
		if(l < 1)
			l = 1;
		if(r > len)
			r = len;
		
		return query(r) - query(l - 1);
	}
	static void update(int targetIdx, int value) {
		
		for(int i=targetIdx; i<=len; i+= i & -i)
			fenwickTree[i] += value;
		
	}
	static int query(int idx) {
		int sum = 0;
		
		for(int i=idx; i>0; i-= i & -i)
			sum += fenwickTree[i];
		
		return sum;
	}
	static class Query implements Comparable<Query>{
		int left, right, idx, factor;
		Query(int l, int r, int i, int f){
			left = l;
			right= r;
			idx = i;
			factor = f;
		}
		@Override
		public int compareTo(Query o) {
			if(factor != o.factor)
				return factor - o.factor;
			
			return (factor & 1) == 0 ? right - o.right : o.right - right;

		}
	}
}