[백준/Java] 2042 구간 합 구하기

박찬병·2024년 11월 26일

Problem Solving

목록 보기
41/48

https://www.acmicpc.net/problem/2042

문제 요약

N개의 정수로 이루어진 수열이 있다. 이때 수열의 한 값이 M번 변경되며, 그 부분합을 K번 구해야 한다.
N은 최대 100만, M과 K는 최대 1만이다.


문제 접근

단순하게 생각하면 값을 변경하고, 구간합을 계산하는 굉장히 간단한 문제로 보인다.
하지만 이렇게 간단하게 문제를 해결하면 구간합을 구하는 횟수 O(K)와 구간합을 구하는 시간복잡도 O(N)에 의해 전체 시간복잡도가 O(KN)이 된다.
그런데 N이 최대 100만이고, K가 최대 1만이기 때문에 이러한 시간복잡도라면 문제를 시간 내에 해결할 수 없다.

문제를 해결할 키는 바로 인덱스 트리이다.
인덱스 트리를 사용하면 특정 구간합을 O(logN)으로 얻을 수 있기 때문에 전체 시간복잡도는 O(KlogN)가 되어 문제를 해결할 수 있다.

주어진 숫자 배열을 인덱스 트리로 구성하고, 값을 변경하며 부분합을 구하면 된다.

인덱스 트리로 부분합 구하기

  1. left가 right child를 가리키면 그 값을 더하고 left++, 아니면 가만히 있음
  2. right가 left child를 가리키면 그 값을 더하고 right--, 아니면 가만히 있음
  3. left와 right가 부모로 이동함
  4. 이를 left와 right가 교차될 때까지 반복 (같아지면 그 값을 더하고 끝냄)

풀이

기본적인 아이디어는 다음과 같다.

  1. (정리 예정)

이를 구현한 코드는 다음과 같다.

import java.util.*;
import java.io.*;

public class Main {
    
    static int N, M, K;
    static long[] numArr;
    static long[][] orderArr;
    
    static long[] numIndexTree;
    static int leafPointer;
    
    // 입력 받기
    public static void getInput() throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        		
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());
        
        // 숫자 배열 얻기
        numArr = new long[N];
        for (int i = 0; i < N; i++) {
        	numArr[i] = Long.parseLong(br.readLine());
        }
        
        // 명령 배열 얻기
        orderArr = new long[M+K][3];
        for (int j = 0; j < M+K; j++) {
        	StringTokenizer st1 = new StringTokenizer(br.readLine());
        	long[] order = new long[3];
        	for (int x = 0; x < 3; x++) {
        		order[x] = Long.parseLong(st1.nextToken());
        	}
        	orderArr[j] = order;
        }
    }
    
    // 인덱스 트리 생성하기
    public static void makeIndexTree() {
    	
    	leafPointer = 1;
    	// N보다 큰, 가장 작은 2의 제곱수 구하기
    	// 지금은 괜찮은데, 나중에는 bit shift 사용 안하면 시간초과되는 문제가 있을 수도 있음
    	while (leafPointer < N) {
    		leafPointer *= 2;
    	}
    	
    	int treeSize = leafPointer * 2;
    	numIndexTree = new long[treeSize];
    	
    	// 트리의 요소를 0으로 초기화
    	for (int i = 0; i < treeSize; i++) {
    		numIndexTree[i] = 0;
    	}
    	
    	// 배열의 값 적기
    	for (int j = 0; j < N; j++) {
    		numIndexTree[leafPointer+j] = numArr[j];
    	}
    	
    	// 합을 구하면서 인덱스 트리 완성하기
    	for (int k = leafPointer-1; k > 0; k--) {
    		numIndexTree[k] = numIndexTree[2*k] + numIndexTree[2*k+1];
    	}
    }
    
    // 인덱스 트리의 특정 값을 변경하기
    // 타겟 인덱스가 1부터 시작하기 때문에 이를 고려해야 함
    public static void updateIndexTree(int targetIndex, long newValue) {
    	targetIndex = targetIndex + leafPointer - 1;
    	
    	numIndexTree[targetIndex] = newValue;
    	targetIndex /= 2;
    	
    	while (targetIndex > 0) {
    		numIndexTree[targetIndex] = numIndexTree[2*targetIndex] + numIndexTree[2*targetIndex+1];
    		targetIndex /= 2;
    	}
    }
    
    // 왼쪽 자식인지 판단하기
    public static boolean isLeftChild(int targetIndex) {
    	if (targetIndex % 2 == 0) {
    		return true;
    	}
    	return false;
    }
        
    // 인덱스 트리의 구간합을 구하기
    // index가 교차되면 끝나야 하고, index가 같다면 해당 값까지 더하고 끝냄
    public static long calPartSum(int leftIdx, int rightIdx) {
    	long partSum = 0;
    	
    	leftIdx = leftIdx + leafPointer - 1;
    	rightIdx = rightIdx + leafPointer - 1;
    	
    	while (rightIdx > leftIdx) {
    		if (!isLeftChild(leftIdx)) {
    			partSum += numIndexTree[leftIdx];
    			leftIdx++;
    		}
    		if (isLeftChild(rightIdx)) {
    			partSum += numIndexTree[rightIdx];
    			rightIdx--;
    		}
    		leftIdx /= 2;
    		rightIdx /= 2;
    	}
    	
    	if (rightIdx == leftIdx) {
    		partSum += numIndexTree[rightIdx];
    	}
    	
    	return partSum;
    }
    
    // 명령을 따라가며 정답을 얻기
    public static void printAnswer() {
    	for (int i = 0; i < M+K; i++) {
    		long[] order = orderArr[i];
    		
    		// 값 변경
    		if (order[0] == 1) {
    			updateIndexTree((int) order[1], order[2]);
    		}
    		// 부분합 얻고 출력
    		else if (order[0] == 2) {
    			long partSum = calPartSum((int) order[1], (int) order[2]);
    			System.out.println(partSum);
    		}
    	}
    }
    
    
    
    public static void main(String[] args) throws IOException {
        getInput();
        
        makeIndexTree();
        
        printAnswer();
    }
}

회고

틀렸던 이유
입력으로 주어지는 숫자가 Long의 범위라서 Long을 사용했는데, 입력을 정수형으로 변경할 때 Integer.parseInt를 사용했다가 틀렸다.

0개의 댓글