[코테] 백준 2042번 : 구간 합 구하기 (java)

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

알고리즘 분류 : 자료 구조, 세그먼트 트리

 

❓문제

---

🔅해석

세그먼트 트리를 이용하는 가장 기본적인 문제이다.

인덱스 트리 생성 -  init

인덱스 트리의 리프 노드에 배열을 담아야 하므로 리프 노드가 시작하는 인덱스 번호를 알아야한다.

리프 노드의 개수를 S로 표현하고 S는 N 이상의 2^n으로 표현 가능한 수여야 한다.

따라서 S 는 아래와 같이 구한다. (N은 주어진 배열의 크기)

S = 1;
while (S < N) {
    S *= 2;
}

 

리프 노드에 배열에 적혀있는 수를 넣기위해 S번째부터 넣는다.

그 외 내부 노드는 S - 1번째부터 자식 노드(왼쪽, 오른쪽) 합을 대입한다.

private static void init() {
    // leaf노드 채우기
    for (int i = 0; i < N; i++) {
        tree[S + i] = nums[i];
    }

    // 내부노드 채우기
    for (int i = S - 1; i > 0; i--) {
        tree[i] = tree[i * 2] + tree[i * 2 + 1];
    }
}

배열로 나타낸 인덱스 트리

 

 

인덱스 트리 query 수행

• 예를 들어 3 ~ 7 구간의 합을 구하는 쿼리를 날린다고 가정한다.

1. 루트 노드에서 시작해서 트리를 전위 순회하는 방식으로 탐색한다.
2. 현재 노드가 구하고자하는 구간에 완전히 속해있으면 그 노드의 결과값을 반환한다.
3. 왼쪽 자식에서 반환된 값과 오른쪽 자식에서 반환된 값을 더해서 반환한다.

• queryLeft와 queryRight는 카운팅 쿼리가 아니면 변하지않고 그대로 넘긴다.

인덱스 트리 query

 

구간 판별은 아래와 같이 그림으로 기억하면 좋다.

query 구간 판별

 

 

인덱스 트리 update 수행

1. 해당 노드 위치의 값을 수정한다.
2. 부모를 따라 올라가면서 결과값을 수정한다.
3. 3번 위치의 데이터가 3으로 바뀐다면 다음 색칠한 노드들이 수정된다.

 

인덱스 트리 update

update 구간 판별

 

---

⭕정답 코드

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

public class Main {

	static int N, M, K;
	static long[] nums;
	static long[] tree;
	static int S;

	public static void main(String[] args) 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());

		S = 1;
		while (S < N) {
			S *= 2;
		}

		nums = new long[N];
		tree = new long[S * 2];

		for (int i = 0; i < N; i++) {
			nums[i] = Long.parseLong(br.readLine());
		}

		init();

		for (int i = 0; i < M + K; i++) {
			st = new StringTokenizer(br.readLine());
			int type = Integer.parseInt(st.nextToken());

			if (type == 1) {
				int idx = Integer.parseInt(st.nextToken());
				long val = Long.parseLong(st.nextToken());
				long diff = val - tree[S + idx - 1];
				update(1, S, 1, idx, diff);
			} else {
				int from = Integer.parseInt(st.nextToken());
				int to = Integer.parseInt(st.nextToken());
				System.out.println(query(1, S, 1, from, to));
			}
		}
		br.close();

	}

	private static void init() {
		// leaf노드 채우기
		for (int i = 0; i < N; i++) {
			tree[S + i] = nums[i];
		}

		// 내부노드 채우기
		for (int i = S - 1; i > 0; i--) {
			tree[i] = tree[i * 2] + tree[i * 2 + 1];
		}
	}

	static long query(int left, int right, int node, int queryLeft, int queryRight) {
		if (queryRight < left || right < queryLeft) { // 1. 상관없는 경우
			return 0;
		} else if (queryLeft <= left && right <= queryRight) { // 2. 판단 가능한 경우
			return tree[node];
		} else { // 3. 판단 불가한 경우(범위에 속하지만 정확하지 않음)
			int mid = (left + right) / 2;
			// 왼쪽, 오른쪽 자식 합한 값을 리턴
			return query(left, mid, node * 2, queryLeft, queryRight)
					+ query(mid + 1, right, node * 2 + 1, queryLeft, queryRight);
		}
	}

	static void update(int left, int right, int node, int target, long diff) {

		if (target < left || right < target) { // 1. 영향이 없는 경우
			return;
		} else { // 2. 영향이 미치는 경우
			tree[node] += diff;
			if (left != right) {
				int mid = (left + right) / 2;
				update(left, mid, node * 2, target, diff);
				update(mid + 1, right, node * 2 + 1, target, diff);
			}
		}
	}
}

 

---

❗결과