백준 Java | 세그먼트 트리 문제 풀이(BOJ 1275번 커피숍2)

1275번: 커피숍2 (acmicpc.net)

1275번: 커피숍2

세그먼트 트리는 3단계로 구성된다. 

1. 트리 만들기

2. 트리 업데이트하기

3. 트리 쿼리(누적합 구하기)

 

처음에 잘못했던 것

- 트리 업데이트에서는 트리 배열 뿐만 아니라 원본 데이터도 값을 업데이트해주어야 한다. 

package SegmentTree.P1275;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.sql.SQLOutput;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {
    static int N;
    static int Q;
    static long[] data;
    static long[] tree;
    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());
        Q = Integer.parseInt(st.nextToken());

        st = new StringTokenizer(br.readLine());
        data = new long[N];
        for (int i=0;i<N;i++){
            data[i] = Long.parseLong(st.nextToken());
        }

        Init();
        for (int i=0;i<Q;i++){
            st = new StringTokenizer(br.readLine());
            int x = Integer.parseInt(st.nextToken())-1;
            int y = Integer.parseInt(st.nextToken())-1;
            // x가 큰 경우 고려해야 함;
            if (x > y){
                System.out.println(Query(y, x));
            } else {
                System.out.println(Query(x, y));
            }
            int a = Integer.parseInt(st.nextToken())-1;
            long b = Long.parseLong(st.nextToken());
            Update(a, b);
        }
    }

    private static void Init() {
        // tree 의 사이즈를 정해요
        int treeSize = 1;
        while (treeSize < N){
            treeSize *= 2;
        }

        tree = new long[treeSize*2+1];

        InitTree(1, 0, N-1);
    }

    private static long InitTree(int node, int start, int end) {
        if (start==end){
            return tree[node] = data[start];
        }
        int mid = (start + end) / 2;
        return tree[node] =
                InitTree(node * 2, start, mid)
                + InitTree(node * 2 + 1, mid + 1, end);
    }

    private static void Update(int a, long b) {
        // a번째 수를 b로 바꾸어라
        UpdateNode(1, 0, N-1, a, b-data[a]);
    }

    private static void UpdateNode(int node, int start, int end, int target, long diff) {
        if (start > target || end < target){
            return;
        } else {
            tree[node] += diff;
            if (start != end){
                int mid = (start + end) / 2;
                UpdateNode(node * 2, start, mid, target, diff);
                UpdateNode(node * 2 + 1, mid + 1, end, target, diff);
            }
            if (start == end){
                data[target] += diff;
            }
        }
    }

    private static long Query(int start, int end) {
        // x번째 수부터 y번째 수까지의 합
        return QuerySum(1, 0, N-1, start, end);
    }

    private static long QuerySum(int node, int start, int end, int targetStart, int targetEnd) {
        // 해당 사항 없음
        // 범위에 쏙 들어감
        // 범위에 걸쳐있음
        if (targetEnd < start || targetStart > end){
            return 0;
        } else if (targetStart <= start && end <= targetEnd){
            return tree[node];
        } else {
            int mid = (start + end) / 2;
            return QuerySum(node * 2, start, mid, targetStart, targetEnd)
                    + QuerySum(node * 2 + 1, mid + 1, end, targetStart, targetEnd);
        }
    }
}