[자바] 백준 16978 - 수열과 쿼리 22 (java)

 문제 : boj16978

---

 

필요 알고리즘 개념

• 오프라인 쿼리
  
  • 온라인 쿼리로 풀려면 어려운 문제지만, 오프라인 쿼리 아이디어가 떠올랐다면 훨씬 쉬워진다.
• 세그먼트 트리 또는 펜윅 트리
  
  • point update range query 문제로 세그먼트 혹은 펜윅 트리를 사용해 풀 수 있다.

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

 

---

 

풀이

  point update와 range query가 존재하는 문제인데, 잘 생각해보면 update와 query를 굳이 실시간으로 해줄 필요가 없다는 점을 알 수 있다. 이게 핵심 아이디어인 문제이다. 즉 이하 로직으로 진행한다.

 

1. '1' 쿼리와 '2'쿼리를 나눠서 별도로 리스트에 담아둔다. 이 때 '2'쿼리의 경우, 해당 쿼리가 나왔던 위치를 별도로 기록해둔다. (코드의 query2Idx)

 

2. '2'쿼리를 k값을 기준으로 정렬한다.

 

3. 그럼 이제 '2'쿼리를 순회하면서, k값이 변경되면 해당 '1'쿼리까지 적용시키면 된다! 만약 '2'쿼리가 k값을 기준으로 정렬되어 있지 않았다면 매번 '1'쿼리를 적용했다가, 제거했다가 해야하지만 k값을 기준으로 정렬해뒀으므로 매번 '1'쿼리를 삽입하기만 하면 된다.

 

  이 때, 쿼리 처리는 세그먼트 트리 혹은 펜윅트리를 사용하면 된다. 누적합 알고리즘의 경우엔 중간의 값이 변경될 경우, 그 이후 모든 값이 변경되어야 하므로 시간초과가 나게 된다. 내 경우엔 펜윅트리로 풀었다. 혹시 펜윅트리에 대해 모른다면 적어둔 '펜윅트리 글'에서 이 문제의 경우엔 '기본' 부분만 보면 풀 수 있다.

 

---

 

코드 : github

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.StringTokenizer;

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

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

    private long pSum(int idx) {
        long answer = 0l;
        while (idx > 0) {
            answer += bit[idx];
            idx -= idx&-idx;
        }
        return answer;
    }

    private long query(int a, int b) {
        return pSum(b) - pSum(a-1);
    }

    private void solution() throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in), 1<<17);
        n = Integer.parseInt(br.readLine());
        arr = new int[n+1];
        bit = new long[n+1];
        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
            update(i, arr[i]);
        }
        int m = Integer.parseInt(br.readLine());
        ArrayList<int[]> query1 = new ArrayList<>();
        query1.add(null);
        ArrayList<int[]> query2 = new ArrayList<>();
        int query2Idx = 0;
        while (m-->0) {
            st = new StringTokenizer(br.readLine());
            int op = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            switch (op) {
                case 1 : query1.add(new int[]{a, b}); break;
                case 2 : query2.add(new int[]{a, b, Integer.parseInt(st.nextToken()), query2Idx++}); break;
            }
        }

        Collections.sort(query2, new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[0]-o2[0];
            }
        });

        int bf = 0;
        long[] answer = new long[query2.size()];
        for (int[] cur : query2) {
            if (bf != cur[0]) {
                for (int i = bf+1; i <= cur[0]; i++) {
                    int[] q = query1.get(i);
                    update(q[0], q[1] - arr[q[0]]);
                    arr[q[0]] = q[1];
                    bf = cur[0];
                }
            }
            answer[cur[3]] = query(cur[1], cur[2]);
        }

        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < answer.length; i++) {
            sb.append(answer[i]).append('\n');
        }
        System.out.print(sb);
    }
    public static void main(String[] args) throws Exception {
        new Main().solution();
    }
}