[BOJ/Segment Tree] 백준 2268 - 수들의 합 7 (Java)

2268 - 수들의 합 7

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

문제

N개의 수 A[1], A[2], ..., A[N]이 주어졌을 때
Sum(i, j) -> A[i] + A[i + 1] + ... + A[j] / Modify(i, k) -> A[i] = k

입력
첫째 줄: N, 수행한 명령의 개수 M (1 <= N, M <= 1,000,000) 
M개 줄: 수행한 순서대로 함수 목록
ㄴ어느 함수 사용했는지(0: Sum 함수/ 1: Modify 함수), 함수의 인자 (i, j) or (i, k)
A[1] = A[2] = ... = A[N] = 0
Modify인 경우에 1 <= k <= 100,000
출력: Sum 함수의 개수만큼 각 줄에 Sum 함수의 리턴값 출력

---

풀이

• 0 a b -> a ~ b 구간 합
• 1 a b -> A[a] = b (값 변경)

 

펜윅 트리(Fenwick Tree)

• `tree`: 내부 트리 / `arr`: 실제 값
• `prefix(i)`: 1 ~ i 까지 누적합
  • `i - (i & -i)`: 현재 구간 크기만큼 왼쪽으로 이동
    • `i & -i`: i의 이진수에서 가장 오른쪽에 있는 1만 남김
      = i가 포함하는 가장 작은 2의 거듭제곱
      -> 현재 노드가 담당하는 구간 크기
    • ex. i = 12 = 00001100/ 1의 보수(모든 비트 뒤집기) = 11110011
      / +1 -> -12 = 11110100 -> i & -i = 00000100
      => 12의 이진수 1100에서 오른쪽에서 처음 나오는 1은 0100 (4)
• `query(lt, rt)`: 누적합 공식 / 시간복잡도: O(log N)
• `update(i, diff)`: i번 값이 `diff`만큼 변했을 때 구간 갱신 / 시간복잡도: O(log N)
  • `i + (i & -i)`: 현재 index를 포함하는 상위 구간으로 이동

static long[] tree, arr;

static long prefix(int i) {
    long sum = 0;
    while (i > 0) {
        sum += tree[i];
        i -= (i & -i);
    }
    return sum;
}

static long query(int lt, int rt) {
    return prefix(rt) - prefix(lt - 1);
}

static void update(int i, long diff) {
    while (i <= N) {
        tree[i] += diff;
        i += (i & -i);
    }
}

 

세그먼트 트리(Segment Tree)

static long[] tree, arr;

static long query(int node, int start, int end, int lt, int rt) {
    if (rt < start || end < lt) return 0;
    if (lt <= start && end <= rt) return tree[node];

    int mid = (start + end) / 2;
    return query(node * 2, start, mid, lt, rt)
            + query(node * 2 + 1, mid + 1, end, lt, rt);
}

static void update(int node, int start, int end, int index, long diff) {
    if (index < start || index > end) return;

    tree[node] += diff;

    if (start != end) {
        int mid = (start + end) / 2;
        update(node * 2, start, mid, index, diff);
        update(node * 2 + 1, mid + 1, end, index, diff);
    }
}

 

*참고

[BOJ/Segment Tree] 백준 2042 - 구간 합 구하기 / 10999 - 구간 합 구하기 2 (Java)

---

코드

펜윅 트리

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

// 수들의 합 7
public class boj_2268 {
    static int N, M;
    static long[] tree, arr;

    static long prefix(int i) {
        long sum = 0;
        while (i > 0) {
            sum += tree[i];
            i -= (i & -i);
        }
        return sum;
    }

    static long query(int lt, int rt) {
        return prefix(rt) - prefix(lt - 1);
    }

    static void update(int i, long diff) {
        while (i <= N) {
            tree[i] += diff;
            i += (i & -i);
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();

        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());

        tree = new long[N + 1];
        arr = new long[N + 1];
        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int type = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            if (type == 0) {
                int lt = Math.min(a, b);
                int rt = Math.max(a, b);
                sb.append(query(lt, rt)).append("\n");
            } else {
                long diff = b - arr[a];
                arr[a] = b;
                update(a, diff);
            }
        }

        System.out.println(sb);
    }
}

 

세그먼트 트리

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

// 수들의 합
public class Main {
    static int N, M;
    static long[] tree, arr;

    static long query(int node, int start, int end, int lt, int rt) {
        if (rt < start || end < lt) return 0;
        if (lt <= start && end <= rt) return tree[node];

        int mid = (start + end) / 2;
        return query(node * 2, start, mid, lt, rt)
                + query(node * 2 + 1, mid + 1, end, lt, rt);
    }

    static void update(int node, int start, int end, int index, long diff) {
        if (index < start || index > end) return;

        tree[node] += diff;

        if (start != end) {
            int mid = (start + end) / 2;
            update(node * 2, start, mid, index, diff);
            update(node * 2 + 1, mid + 1, end, index, diff);
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();

        N = Integer.parseInt(st.nextToken());
        M = Integer.parseInt(st.nextToken());

        tree = new long[4 * N];
        arr = new long[N + 1];

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

            if (type == 0) {
                int lt = Math.min(a, b);
                int rt = Math.max(a, b);
                sb.append(query(1, 1, N, lt, rt)).append("\n");
            } else {
                long diff = b - arr[a];
                arr[a] = b;
                update(1, 1, N, a, diff);
            }
        }

        System.out.print(sb);
    }
}