BOJ 백준 8462 배열의 힘

 

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

[ 문제 요약 ]

• 범위가 주어지면, 그 범위 안에 있는 특정수 S가 몇 개 있는지를 세고, 그 숫자를 K라 하면 K * K * S를 구하고, 모든 특정수 S에 대해 이 연산을 반복해 모두 더한 후 출력합니다.

[ 테스트 케이스 설명 ]

8 3			// 배열 크기(1<=100,000), 부분 배열 개수(1<=100,000)
4 3 1 1 1 3 1 2		// 1<=1,000,000
2 7			// 범위
1 6
3 8
답 / 범위 안에있는 특정수 S가 몇개있는지를 세고, 그 숫자를 K라하면 K * K * S를 출력
28
25
21

 

[ 문제 해설 ]

쿼리들을 먼저 받아 오프라인 쿼리로 처리합니다.

 

그리고 빠른 투포인터 연산을 위해 Mo’s 알고리즘으로 쿼리를 정렬한 후 풀면 됩니다.

 

Mo’s 알고리즘은 배열을 sqrt(N) 구간으로 나눠 최소한의 움직임으로 쿼리 처리를 할 수 있도록 입력된 쿼리들을 정렬하는 데 사용됩니다.

 

쿼리 정렬 후 단순히 투 포인터 문제를 풀듯이 풀면 됩니다.

 

정답이 int를 넘어가기 때문에 long으로 선언해야 합니다. 그리고 특정수 S를 카운팅 할 때마다 최종 결과에 값을 추가하고 삭제하는 일을 반복해야 합니다. 그래야 숫자 카운팅 시마다 결과를 바로바로 알 수 있습니다.

 

그래서 정답 코드를 보면 +와 -를 번갈아 가면서 진행합니다. +,-를 번갈아 사용하지 않고, 수학적인 계산으로 한 번에 연산 가능하도록 코드를 수정할 수도 있습니다.

 

아래는 일반적인 정답 코드와 중복 연산을 줄인 코드를 보여줍니다.

 

[ 정답 코드 ]

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.StringTokenizer;
class Main{
	
	static int sqrt;
	
	static class Query implements Comparable<Query>{
		int left, right, idx;
		Query(int l, int r, int i){
			left = l;
			right = r;
			idx = i;
		}
		@Override
		public int compareTo(Query o) {
			int l = left / sqrt;
			int r = o.left / sqrt;
			return l == r ? right - o.right : l - r;
		}
	}
	
	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 Q		= Integer.parseInt(st.nextToken());// 부분 배열 개수(1<=100,000)
		int arr[]	= new int[N + 1];
		long cnt[]	= new long[1_000_001];
		long ans[]	= new long[Q + 1];
		sqrt		= (int)Math.sqrt(N);
		
		st = new StringTokenizer(br.readLine());
		for(int i=1; i<=N; i++)
			arr[i] = Integer.parseInt(st.nextToken());// 1<=1,000,000
		
		ArrayList<Query> 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());
			query.add(new Query(l, r, i));
		}
		
		Collections.sort(query);
		
		long now = 0;
		int idxL = 1;
		int idxR = 0;
		for(Query q : query)
		{			
			while(idxR < q.right)
			{
				++idxR;
				int num = arr[idxR];
				
				now -= cnt[num] * cnt[num] * num;
				++cnt[num];
				now += cnt[num] * cnt[num] * num;
			}
			while(q.right < idxR)
			{
				int num = arr[idxR];
				
				now -= cnt[num] * cnt[num] * num;
				--cnt[num];
				now += cnt[num] * cnt[num] * num;
				
				--idxR;
			}
			while(idxL < q.left)
			{
				int num = arr[idxL];
				
				now -= cnt[num] * cnt[num] * num;
				--cnt[num];
				now += cnt[num] * cnt[num] * num;
				
				++idxL;
			}
			while(q.left < idxL)
			{
				--idxL;
				int num = arr[idxL];
				
				now -= cnt[num] * cnt[num] * num;
				++cnt[num];
				now += cnt[num] * cnt[num] * num;
			}
			
			ans[q.idx] = now; 
		}
		
		StringBuilder sb = new StringBuilder();
		for(int i=1; i<=Q; i++)
			sb.append(ans[i]).append('\n');
		
		System.out.print(sb);
	}
}

 

 

[ 한번에 연산하는 코드 ]

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.StringTokenizer;
class Main{
	
	static int sqrt;
	
	static class Query implements Comparable<Query>{
		int left, right, idx;
		Query(int l, int r, int i){
			left = l;
			right = r;
			idx = i;
		}
		@Override
		public int compareTo(Query o) {
			int l = left / sqrt;
			int r = o.left / sqrt;
			return l == r ? right - o.right : l - r;
		}
	}
	
	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 Q		= Integer.parseInt(st.nextToken());// 부분 배열 개수(1<=100,000)
		int arr[]	= new int[N + 1];
		long cnt[]	= new long[1_000_001];
		long ans[]	= new long[Q + 1];
		sqrt		= (int)Math.sqrt(N);
		
		st = new StringTokenizer(br.readLine());
		for(int i=1; i<=N; i++)
			arr[i] = Integer.parseInt(st.nextToken());// 1<=1,000,000
		
		ArrayList<Query> 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());
			query.add(new Query(l, r, i));
		}
		
		Collections.sort(query);
		
		long now = 0;
		int idxL = 1;
		int idxR = 0;
		for(Query q : query)
		{			
			while(idxR < q.right)
			{
				++idxR;
				
				int num = arr[idxR];
				
				now += 2*cnt[num] * num + num;
				
				++cnt[num];
			}
			while(q.right < idxR)
			{
				int num = arr[idxR];
				
				now -= 2*cnt[num] * num - num;
				
				--cnt[num];
				
				--idxR;
			}
			while(idxL < q.left)
			{
				int num = arr[idxL];
				
				now -= 2*cnt[num] * num - num;
				
				--cnt[num];
				
				++idxL;
			}
			while(q.left < idxL)
			{
				--idxL;
				
				int num = arr[idxL];

				now += 2*cnt[num] * num + num;
				
				++cnt[num];
			}
			
			ans[q.idx] = now; 
		}
		
		StringBuilder sb = new StringBuilder();
		for(int i=1; i<=Q; i++)
			sb.append(ans[i]).append('\n');
		
		System.out.print(sb);
	}
}

 

한 번에 연산하는 코드가 그나마 속도가 빠른 것을 볼 수 있습니다. 위와 같은 식이 나온 것은 아래 식들을 Z에 대해서 정리하면 나오게 됩니다. Z는 현재 얼마를 더해야 둘이 같아지는지를 의미합니다. X는 특정수 S가 몇 번 있는지를 의미, Y는 특정수 S를 의미합니다.

 

특정수 S의 카운팅이 늘어날 때 : X * X * Y + Z = (X+1) * (X+1) * Y

특정수 S의 카운팅이 줄어들 때 : X * X * Y - Z = (X-1) * (X-1) * Y