BOJ 백준 13554 수열과 쿼리 9

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

 

[ 문제 요약 ]

• 배열 A, B가 주어지고, 구간 쿼리가 주어질 때 아래 질문의 답을 출력합니다.
• i j k : i ≤ p, q ≤ j이면서 Ap × Bq ≤ k인 (p, q) 쌍의 개수를 출력한다.

 

[ 테스트 케이스 설명 ]

2		// 수열의 크기(1<=100,000)
1 3	// 배열(1) 1<=100,000
2 4	// 배열(2) 1<=100,000
4		// 쿼리 개수(1<=100,000)
1 2 12// i, j, k(1<=100,000)가 주어짐
1 1 5
1 2 6
2 2 4
정답 : i j k : i ≤ p, q ≤ j이면서 Ap × Bq ≤ k 인 (p, q) 쌍의 개수를 출력한다.
4
1
3
0

 

 

[ 알고리즘 분류 ]

• 자료 구조
• 세그먼트 트리
• 오프라인 쿼리
• 제곱근 분할법
• mo’s

 

[ 문제 해설 ]

 

mo’s 문제 대부분은 구간에 대해 mo’s 알고리즘으로 정렬하고 투 포인터를 이동하며 답을 구하고, 과거의 답을 현재에 활용할 수 있는 게 대부분입니다.

 

하지만 이 문제는 쿼리마다 K가 주어지기 때문에 기존 답을 활용할 수 없어 항상 다시 계산해야 합니다.

 

A * B ≤ K 일 때, 아래 식과 같이 A를 알면 B의 범위를 알 수 있습니다. 역으로 B를 알면 A의 범위를 구할 수 있습니다.

 

B를 알 때 A의 범위 : A ≤ K / B

A를 알 때 B의 범위 : B ≤ K / A

 

이렇게 되면, 세그먼트 트리의 인덱스를 각 배열의 인자 값들로 하고, 세그먼트 트리에 저장되는 value를 각 배열 인자들의 등장한 횟수로 해놓으면 A나 B를 알 때 가능한 범위를 구해 상대 배열에서 곱했을 때 K 이하인 쌍을 빠르게 알 수 있습니다.

 

그런데 한 배열을 기준으로 다른 배열을 완전 탐색한다면 시간 초과입니다.

 

효율적인 탐색을 위해 쿼리마다 주어지는 K의 제곱근을 활용합니다.

 

K 제곱근을 활용하는 이유는, 두 수의 곱에서 둘 중 하나는 반드시 K의 제곱근 보다 작거나 같습니다.

 

이를 활용해 1부터 K의 제곱근까지 순회하면서 구간에 대한 질의를 빠르게 구해줍니다.

 

예를 들어 K가 16으로 양 배열에서 곱해서 16이 되는 쌍을 빨리 찾아야 하면,

 

K의 제곱근인 4를 통해 반복문으로 구해줍니다.

 

아래는 코드 예시입니다.

for(int i = 1; i<= sq; i++)// sq는 K의 제곱근
{
	// 중복을 제거하며 모두 탐색해야 하기 때문에 B를 기준으로 나눔
	// 첫번 째 B는 1 ~ sq까지				/	A는 1 ~ 가능한 부분까지
	// 두번 째 B는 sq + 1 ~ 가능한 부분까지	/	A는 1 ~ sq까지
	res += query(i, i, fenwickTree_B) * query(1, q.k/i, fenwickTree_A);
	res += query(sq + 1, q.k/i, fenwickTree_B) * query(i, i, fenwickTree_A);
}

 

A[] = { 1 , 2, 3, 4, 5 }

B[] = { 2, 3, 4, 5, 6 }

 

위와 같은 배열이 있을 때, B를 기준으로 나눠 처리합니다.(A로 구분해도 동일) 나누는 이유는 중복을 제거하기 위함입니다.

 

sq 값이 4일 때, 배열 B[] 입장에서 위 코드를 보면,

 

첫 번째 연산에서 B[]는 4 이하의 2,3,4를 순차 탐색하고,

두 번째 연산에서 B[]는 4초과의 범위를 탐색하게 됩니다.

 

그리고 그에 따른 A 배열을 상황에 맞게 대입해 줍니다. B가 하나로 고정일 때는 A 배열이 범위 스캔, B가 범위 스캔일 때는 A를 고정해서 탐색합니다.

 

이렇게 하면 sqrt(K) 횟수만에 쌍을 구할 수 있습니다.

 

추가로 세그먼트 트리는 update와 query 연산시 재귀 오버헤드가 있어 느립니다.

 

단순 구간 합, 구간 업데이트는 펜윅 트리로 하는 것이 빠르기 때문에 저는 펜윅 트리로 구현하였습니다.

 

[ 정답 코드 ]

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.StringTokenizer;
class Main{
	
	static final int LEN = 100_000;
	static int N, Q;
	static int sqrt;
	static int[] A, B;
	static int[] fenwickTree_A, fenwickTree_B;
	static long[] ans;
	static long[] diff;
	static ArrayList<Query> query;
	
	public static void main(String[] args)throws Exception{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		N = Integer.parseInt(br.readLine());
		A = new int[N + 1];
		B = new int[N + 1];
		diff = new long[N + 1];
		sqrt = (int)Math.sqrt(N);
		fenwickTree_A = new int[LEN + 1];
		fenwickTree_B = new int[LEN + 1];
		
		StringTokenizer st = new StringTokenizer(br.readLine());
		for(int i=1; i<=N; i++)
			A[i] = Integer.parseInt(st.nextToken());
		
		st = new StringTokenizer(br.readLine());
		for(int i=1; i<=N; i++)
			B[i] = Integer.parseInt(st.nextToken());
		
		Q = Integer.parseInt(br.readLine());
		ans = new long[Q + 1];
		query = new ArrayList<>();
		
		for(int i=1; i<=Q; i++)
		{
			st = new StringTokenizer(br.readLine());
			int l = Integer.parseInt(st.nextToken());
			int r = Integer.parseInt(st.nextToken());
			int k = Integer.parseInt(st.nextToken());
			query.add(new Query(l, r, k, i, l/sqrt));
		}
		
		Collections.sort(query);
		
		int L = 1;
		int R = 0;

		for(Query q : query)
		{
			while(R < q.right) excute(++R, 1);
			while(q.left < L) excute(--L, 1);
			while(q.right < R) excute(R--, -1);
			while(L < q.left) excute(L++, -1);

			int sq = (int)Math.sqrt(q.k);

			long res = 0;

			for(int i = 1; i<= sq; i++)
			{
				// 중복을 제거하며 모두 탐색해야 하기 때문에 B를 기준으로 나눔
				// 첫번 째 B는 1 ~ sq까지				/	A는 1 ~ 가능한 부분까지
				// 두번 째 B는 sq + 1 ~ 가능한 부분까지	/	A는 1 ~ sq까지
				res += query(i, i, fenwickTree_B) * query(1, q.k/i, fenwickTree_A);
				res += query(sq + 1, q.k/i, fenwickTree_B) * query(i, i, fenwickTree_A);
			}

			ans[q.idx] = res;
		}
		
		print();
	}
	static void excute(int idx, int plus) {
		update(A[idx], plus, fenwickTree_A);
		update(B[idx], plus, fenwickTree_B);
	}
	static void update(int idx, int plus, int[] tree) {
		while(idx <= LEN)
		{
			tree[idx] += plus;
			idx += idx & - idx;
		}
	}
	static long query(int l, int r, int[] tree) {
		return sum(r, tree) - sum(l - 1, tree);
	}
	static long sum(int idx, int[] tree) {
		long res = 0;
		while(idx > 0)
		{
			res += tree[idx];
			idx -= idx & -idx;
		}
		return res;
	}
	static void print() {
		StringBuilder sb = new StringBuilder();
		
		for(int i=1; i<=Q; i++)
			sb.append(ans[i]).append('\n');
		
		System.out.print(sb);
	}
	static class Query implements Comparable<Query>{
		int left, right, k, idx, fac;
		Query(int l, int r, int k, int i, int f){
			this.left = l;
			this.right = r;
			this.idx = i;
			this.k = k;
			this.fac = f;
		}
		@Override
		public int compareTo(Query o) {
			if(fac != o.fac) // 구간이 다르면 구간 기준 오름차순
				return fac - o.fac;
			// 구간이 짝수인 경우 right 기준 오름차순
			if((fac & 1)==0)
				return right - o.right;
			// 구간이 홀수인 경우 right 기준 내림차순
			return o.right - right;
		}
	}
}