[PS] BOJ2042 구간 합 구하기 ( prefixSum ) With JAVA, Python

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

2042번: 구간 합 구하기

 

 

◎문제풀이

 

[Alogorithm] 세그먼트 트리( Segment Tree ) 구간합 구하기 ( + BOJ2042 ) With Python,JAVA

 

 

◎코드

 

Python

#BOJ2042 구간 합 구하기
import sys
input = sys.stdin.readline
 
n,m,k = map(int,input().split())
arr = [0]*(n+1) # 수열
for i in range(1,n+1) :
  arr[i] = int(input())
tree = [0]*(4*n) # 세그먼트 트리
 
#세그먼트 트리 초기화
def init(node,start,end) :
  if start == end : #리프노드인 경우
    tree[node] = arr[start]
  else :  #리프노드가 아닌 경우
    mid = ( start + end ) // 2
    tree[node] = init(2*node,start,mid) + init(2*node+1,mid+1,end)
  return tree[node]
 
#세그먼트 트리 구간합 탐색하기
def sum_tree(node,start,end,fl,fr) :
  if fl > end or fr < start : return 0
  if fl <= start and fr >= end : return tree[node]
  mid = ( start + end ) // 2
  return sum_tree(2*node,start,mid,fl,fr) + sum_tree(2*node+1,mid+1,end,fl,fr)
 
#세그먼트 트리 UPDATE 하기
def update(node,start,end,idx,diff) :
  if start > idx or end < idx : return 
  tree[node] += diff
  if start != end :
    mid = (start + end) // 2
    update(2*node,start,mid,idx,diff)
    update(2*node+1,mid+1,end,idx,diff)
 
#문제풀이 시작
init(1,0,len(arr)-1) # 세그먼트 트리 초기화
for _ in range(m+k) :
  a,b,c = map(int,input().split())
  # 수정
  if a == 1 :
    diff = c - arr[b] # diff, 변경정도
    arr[b] = c # 수열에 변경내용 반영
    update(1,0,len(arr)-1,b,diff)  # 세그먼트리 변경내용 반영
  #구간합출력
  elif a == 2 :
    print(sum_tree(1,0,len(arr)-1,b,c)) # 구간합 탐색 및 출력

 

 

JAVA

import java.io.*;
import java.util.StringTokenizer;
 
//BOJ2042 구간 합 구하기
public class Main {
    static long[] arr;
    static long[] tree;
    static StringTokenizer st;
 
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringBuilder sb = new StringBuilder();
 
        st = new StringTokenizer(br.readLine());
 
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());
        arr = new long[n+1];
        tree = new long[4*n+1];
 
        for(int i =1;i<n+1;i++){
            st = new StringTokenizer(br.readLine());
            arr[i] = Long.parseLong(st.nextToken());
        }
 
        init_tree(1,n,1);
        int z = m+k;
        while(z-- > 0){
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            long c = Long.parseLong(st.nextToken());
 
            //수정
            if(a==1){
                long diff = c - arr[b];
                arr[b] = c;
                update_tree(1,n,1,b,diff);
            }
            //구간합
            else if(a==2){
                sb.append(sum_tree(1,n,1,b,(int)c)).append("\n");
            }
        }
 
        bw.write(sb.toString());
        bw.flush();
        bw.close();
        br.close();
    }
    public static long init_tree(int start,int end,int node){
        if( start == end ){
            tree[node] = arr[start];
        }
        else {
            int mid = (start + end) / 2;
            tree[node] = init_tree(start, mid, 2 * node) + init_tree(mid + 1, end, 2 * node + 1);
        }
        return tree[node];
    }
 
    public static void update_tree(int start,int end,int node,int idx,long diff){
        if(idx < start || end < idx) return;
        tree[node] += diff;
        if ( start != end ){
            int mid = (start+end)/2;
            update_tree(start,mid,2*node,idx,diff);
            update_tree(mid+1,end,2*node+1,idx,diff);
        }
    }
 
    public static long sum_tree(int start,int end, int node,int fs, int fe){
        if ( fe < start || end < fs ) return 0;
        if ( fs <= start && end <= fe ) return tree[node];
 
        int mid = (start+end)/2;
        return sum_tree(start,mid,2*node,fs,fe) + sum_tree(mid+1,end,2*node+1,fs,fe);
    }
}