Baekjoon 2042 구간 합 구하기 - 세그먼트 트리(구간 합, 수정)

 

더 깔끔한 노션을 원한다면

Baekjoon 2042 구간 합 구하기 - 세그먼트 트리(구간 합, 수정)

2042번: 구간 합 구하기

세그먼트 트리

최솟값 문제에서 배운 세그먼트 트리를 응용하여 변경 및 구간 합 문제를 풀이하였다.

최솟값을 구하는 문제와 달리 구간합에서는 init과 query가 변경되어야 했고 update가 추가됐다.

수의 범위가 굉장히 크므로 long 자료형을 써야한다.

init 함수

static long init(long arr[], int left, int right, int node) {
  if (left == right) {
    return segmentArr[node] = arr[left];
  }
  int mid = (left + right) / 2;
  return segmentArr[node] = (init(arr, left, mid, node * 2) + init(arr, mid + 1, right, node * 2 + 1));
}

Math.min이었던 부분을 그냥 더해주면 된다❗

query 함수

static long query(int start, int end, int node, int nodeLeft, int nodeRight) {
  if (start > nodeRight || end < nodeLeft) {
    return 0L;
  }
  if (start <= nodeLeft && end >= nodeRight) {
    return segmentArr[node];
  }
  int mid = (nodeLeft + nodeRight) / 2;
  return query(start, end, node * 2, nodeLeft, mid) + query(start, end, node * 2 + 1, mid + 1, nodeRight);
}

Math.min이었던 부분을 그냥 더해주면 된다❗

범위에 벗어난다면 Intger.Max_VALUE가 아닌 0을 출력해주면 된다.

update 함수

static void update(int start, int end, int node, int index, long diff) {
  if (!(start <= index && index <= end)) {
    return;
  }
  segmentArr[node] += diff;
  int mid = (start + end) / 2;
  if (start != end) {
    update(start, mid, node * 2, index, diff);
    update(mid + 1, end, node * 2 + 1, index, diff);
  }
}

update 함수는 범위에 벗어나면 return, 벗어나지 않는다면 바꾸려는 값을 segmentArr[node]에 더해주면 된다.

start와 end가 같다면 return을 해준다.

전체 코드

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

public class Main {
  public static void main(String[] args) throws IOException {
    P2042();
  }

  static void P2042() throws IOException {
    BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    StringTokenizer st = new StringTokenizer(br.readLine());
    int N, M, K;
    N = Integer.parseInt(st.nextToken());
    M = Integer.parseInt(st.nextToken());
    K = Integer.parseInt(st.nextToken());
    long[] numbers = new long[N];
    for (int index = 0; index < N; index++) {
      numbers[index] = Long.parseLong(br.readLine());
    }
    SegmentTree segObj = new SegmentTree(numbers, N);
    for (int i = 0; i < M + K; i++) {
      st = new StringTokenizer(br.readLine());
      int mode = Integer.parseInt(st.nextToken());
      int a = Integer.parseInt(st.nextToken());
      long b = Long.parseLong(st.nextToken());
      if (mode == 1) { // update
        segObj.update(0, N - 1, 1, a - 1, b - numbers[a - 1]);
        numbers[a - 1] = b;
      } else { // 구간 합
        System.out.println(segObj.query(a - 1, (int) b - 1, 1, 0, N - 1));
      }
    }
  }

  static class SegmentTree {
    static long[] segmentArr;

    SegmentTree(long[] arr, int n) {
      int x = (int) Math.ceil(Math.log(n) / Math.log(2));
      int size = (int) Math.pow(2, x) * 2;
      segmentArr = new long[size];
      Arrays.fill(segmentArr, Long.MAX_VALUE);
      init(arr, 0, n - 1, 1);
    }

    static long init(long arr[], int left, int right, int node) {
      if (left == right) {
        return segmentArr[node] = arr[left];
      }
      int mid = (left + right) / 2;
      return segmentArr[node] = (init(arr, left, mid, node * 2) + init(arr, mid + 1, right, node * 2 + 1));
    }

    static void update(int start, int end, int node, int index, long diff) {
      if (!(start <= index && index <= end)) {
        return;
      }
      segmentArr[node] += diff;
      int mid = (start + end) / 2;
      if (start != end) {
        update(start, mid, node * 2, index, diff);
        update(mid + 1, end, node * 2 + 1, index, diff);
      }
    }

    static long query(int start, int end, int node, int nodeLeft, int nodeRight) {
      if (start > nodeRight || end < nodeLeft) {
        return 0L;
      }
      if (start <= nodeLeft && end >= nodeRight) {
        return segmentArr[node];
      }
      int mid = (nodeLeft + nodeRight) / 2;
      return query(start, end, node * 2, nodeLeft, mid) + query(start, end, node * 2 + 1, mid + 1, nodeRight);
    }
  }
}