[JAVA] 백준 11505 세그먼트 트리

 

 

 

 

11505번: 구간 곱 구하기

세그먼트 트리 개념

[Algorithm] 세그먼트 트리란?

문제

어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 곱을 구하려 한다. 만약에 1, 2, 3, 4, 5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 곱을 구하라고 한다면 240을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 곱을 구하라고 한다면 48이 될 것이다.

이때, 출력을 1,000,000,007로 나눈 나머지를 출력하여라.

풀이

주어진 수들의 변경이 빈번히 일어날 때, 구간에 대한 곱을 구하는 문제이다.

변경이 빈번히 일어나기 때문에 세그먼트 트리를 구성하여 계산하였다.

📌 세그먼트 트리 구성

        int k = 0;
        while (Math.pow(2, k) < N){
            k++;
        }

        int startIndex = (int) Math.pow(2, k);
        tree = new long[startIndex * 2];
        Arrays.fill(tree, 1);

        for(int i=startIndex; i<startIndex+ N; i++){
            tree[i] = Integer.parseInt(br.readLine());
        }
        for(int i=startIndex-1; i>0; i--){
            tree[i] = (tree[i*2] * tree[i*2 +1]) % divide;
        }

$2^k ≥ N$을 만족하는 k의 최솟값을 구하고 k를 토대로 트리 배열의 크기를 지정해주었다.

여기서 트리를 1로 채우는 이유는 곱셈은 0이 있으면 0이 되어버리기 때문이다. (트리에서 사용하지 않는 부분이 0으로 채워져 있으면 안된다.)

• 또 하나 중요한 사항은 리프 노드들을 채운 후 부모 노드들을 채울 때, divide를 미리 해주어서 오버 플로가 발생하지 않도록 해주어야 한다. (나머지 연산)

📌 세그먼트 트리 값 갱신 및 구간 곱 계산

    public static void changeValue(int index, int value){

        tree[index] = value;
        index /= 2;

        while (index > 0){
            tree[index] = (tree[2 * index] * tree[2 * index + 1]) % divide ;
            index /= 2;
        }

    }

    public static long rangeTimes(int startIndex, int endIndex){
        long value = 1;

        while (startIndex <= endIndex){
            if(startIndex % 2 == 1) value = (value * tree[startIndex]) % divide;
            if(endIndex % 2 == 0) value = (value * tree[endIndex]) % divide;

            startIndex = (startIndex + 1) / 2;
            endIndex = (endIndex - 1) / 2;
        }

        return value;
    }

구간 곱은 수가 크기 때문에 항상 곱셈 연산을 하면서 나머지 연산을 같이 해주어야 했다.

📌 출력

        StringBuilder sb = new StringBuilder();
        for(int i=0; i<M + K; i++){
            st = new StringTokenizer(br.readLine(), " ");
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken()) + startIndex-1;
            int c = Integer.parseInt(st.nextToken());
            if(a == 1){
                changeValue(b, c);
            }else{
                c += startIndex-1;
                sb.append(rangeTimes(b, c)).append("\n");
            }

        }
        System.out.println(sb);

📌 전체 코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

class Main {
    static long[] tree;
    static int divide = 1000000007;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine(), " ");

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

        int k = 0;
        while (Math.pow(2, k) < N){
            k++;
        }

        int startIndex = (int) Math.pow(2, k);
        tree = new long[startIndex * 2];
        Arrays.fill(tree, 1);

        for(int i=startIndex; i<startIndex+ N; i++){
            tree[i] = Integer.parseInt(br.readLine());
        }
        for(int i=startIndex-1; i>0; i--){
            tree[i] = (tree[i*2] * tree[i*2 +1]) % divide;
        }

        StringBuilder sb = new StringBuilder();
        for(int i=0; i<M + K; i++){
            st = new StringTokenizer(br.readLine(), " ");
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken()) + startIndex-1;
            int c = Integer.parseInt(st.nextToken());
            if(a == 1){
                changeValue(b, c);
            }else{
                c += startIndex-1;
                sb.append(rangeTimes(b, c)).append("\n");
            }

        }
        System.out.println(sb);

    }
    public static void changeValue(int index, int value){

        tree[index] = value;
        index /= 2;

        while (index > 0){
            tree[index] = (tree[2 * index] * tree[2 * index + 1]) % divide ;
            index /= 2;
        }

    }

    public static long rangeTimes(int startIndex, int endIndex){
        long value = 1;

        while (startIndex <= endIndex){
            if(startIndex % 2 == 1) value = (value * tree[startIndex]) % divide;
            if(endIndex % 2 == 0) value = (value * tree[endIndex]) % divide;

            startIndex = (startIndex + 1) / 2;
            endIndex = (endIndex - 1) / 2;
        }

        return value;
    }

}

결과