본문 바로가기
PS/BOJ

[자바, C++] 백준 16975 - 수열과 쿼리 21 (java cpp)

by Nahwasa 2022. 8. 15.

 문제 : boj16975


 

필요 알고리즘 개념

  •  lazy propagation을 적용한 세그먼트 트리 혹은 펜윅 트리
    • 세그먼트 트리 + lazy propagation 혹은 range update가 가능한 펜윅 트리를 알고 있어야 풀 수 있다.

※ 제 코드에서 왜 main 함수에 로직을 직접 작성하지 않았는지, 왜 Scanner를 쓰지 않고 BufferedReader를 사용했는지 등에 대해서는 '자바로 백준 풀 때의 팁 및 주의점' 글을 참고해주세요. 백준을 자바로 풀어보려고 시작하시는 분이나, 백준에서 자바로 풀 때의 팁을 원하시는 분들도 보시는걸 추천드립니다.

 


 

풀이

  세그먼트 트리 lazy propagation 혹은 range update 펜윅 트리의 기본형태 문제이다. 펜윅 트리로 풀려면 작성해둔 펜윅 트리 글에서 '응용 2' 부분을 읽어보면 이 문제를 풀 수 있다.

 

 


 

코드(C++) : github

#include <iostream>
using namespace std;
#define ll long long

int n;
ll bit[100001] = {0,};

void update(int i, ll diff) {
    while (i <= n) {
        bit[i] += diff;
        i += i&-i;
    }
}

ll query(int i) {
    ll sum = 0;
    while (i > 0) {
        sum += bit[i];
        i -= i&-i;
    }
    return sum;
}

void rangeUpdate(int i, int j, ll diff) {
    update(i, diff);
    update(j+1, -diff);
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);

    int m;
    cin >> n;
    for (int i = 1; i <= n; i++) {
        ll cur;
        cin >> cur;
        rangeUpdate(i, i, cur);
    }

    cin >> m;
    while (m--) {
        int op, a, b, c;
        cin >> op;
        switch (op) {
            case 1:
                cin >> a >> b >> c;
                rangeUpdate(a, b, c);
                break;
            case 2:
                cin >> a;
                cout << query(a) << '\n';
                break;
        }
    }
}

 

 

코드(Java) : github

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

public class Main {
    int n;
    long[] bit;

    private void update(int i, long diff) {
        while (i <= n) {
            bit[i] += diff;
            i += i&-i;
        }
    }

    private long query(int i) {
        long sum = 0;
        while (i > 0) {
            sum += bit[i];
            i -= i&-i;
        }
        return sum;
    }

    private void rangeUpdate(int i, int j, long diff) {
        update(i, diff);
        update(j+1, -diff);
    }

    private void solution() throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        n = Integer.parseInt(br.readLine());
        bit = new long[n+1];
        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            rangeUpdate(i, i, Integer.parseInt(st.nextToken()));
        }
        int m = Integer.parseInt(br.readLine());
        StringBuilder sb = new StringBuilder();
        while (m-->0) {
            st = new StringTokenizer(br.readLine());
            switch (Integer.parseInt(st.nextToken())) {
                case 1: rangeUpdate(Integer.parseInt(st.nextToken()), Integer.parseInt(st.nextToken()), Integer.parseInt(st.nextToken())); break;
                case 2: sb.append(query(Integer.parseInt(st.nextToken()))).append('\n'); break;
            }
        }
        System.out.print(sb);
    }

    public static void main(String[] args) throws Exception {
        new Main().solution();
    }
}

댓글