세그먼트 트리 응용단계 - 이분탐색과 콜라보1

1. 문제

 

1321번: 군인 (acmicpc.net)

1321번: 군인

 

2. 풀이

 

문제를 잘 읽어보면 어떻게 풀어야할지 대충 보인다

 

군인은 군번대로 1번부대부터 차례대로 배치되고, 부대가 4, 3, 7 이렇게 있을때, 군번이 6번인 군인은 2번부대에 배치된다..

 

왜냐하면 1번부대에 4번까지 배치하고 5번,6번 군인은 2번 부대에 배치되니까

 

그런데 부대의 인원이 줄어들거나 늘어날수 있고 그럴때 군인들이 처음부터 다시 배치한다는 것이다

 

이거는 이제 세그먼트 트리의 업데이트과정이지..

 

그러면 이제 query는 리스트의 원소가 변경될때, i번 군인이 어디 부대에 배치되는지를 매번 찾아야하는 문제인데

 

i번 군인이 어디에 배치될지는 어디에 알수 있는가?

 

그냥 구간합을 이용하면 되잖아

 

"4,3,7일때 [0,1]을 합하면 7인데.. 이 7이 배치해야할 6번 군인보다 작으니까 1번 부대에 배치될수 있다" 

 

그러면 구간합을 구할려면 결국 left와 right를 알아야하는데..

 

군인을 처음부터 차례대로 배치하니까 left = 0으로 항상 고정이고

 

결국 right를 움직여가면서 [0,right]의 합이 i번 군인의 군번 i보다 작은지 큰지를 검사해서.. right를 binary search로 찾는다

 

주어진 부대 배열의 크기 n에 대해 start = 0이고 end = n으로 하고..

 

[0,right]의 합이 i보다 크거나 같다면?

 

i번 군인은 right번 부대에 배치될 수 있으므로 end = mid로 하고

 

i보다 작다면..?

 

i번 군인이 배치해야할 부대는 right 기준으로 큰 인덱스쪽에 있으므로, start = mid + 1로 한다..

 

최종적으로 배치해야할 부대는 end번에 있는데.. 부대 번호는 1번부터 세니까 end + 1을 return해야함

 

def binary_search(tree,target,start,end):
    
    while start < end:
        
        mid = start + (end-start)//2

        k = query(tree,n,0,mid+1) #[0,mid]까지 구간합
        
        #구간합이 배치할 군인의 군번 target보다 크거나 같다면? mid에 배치가능
        if k >= target: 
            
            end = mid
        
        #target보다 작다면... 배치해야할 부대 번호는 mid + 1 이후에 있다
        else:
            
            start = mid + 1
    
    #index는 0번부터 세는데.. 부대번호는 1번부터 세야하니까..
    return end+1

 

세그먼트 트리를 만드는 함수들은 기본형태랑 동일해서..

 

binary search를 잘 짜는게 중요하다..

 

그러면 반복문을 이용한 세그먼트 트리와 binary search를 구현하면..

 

from sys import stdin

def create_segment(A,tree,n):
    
    for i in range(n):
        
        tree[n+i] = A[i]
    
    for i in range(n-1,0,-1):
        
        tree[i] = tree[2*i]+tree[2*i+1]

def update(tree,n,index,value):
    
    index += n

    tree[index] += value

    while index > 1:
        
        index = index >> 1

        tree[index] = tree[2*index] + tree[2*index+1]

def query(tree,n,left,right):
    
    left += n
    right += n

    m = 0

    while left < right:
        
        if left & 1:
            
            m += tree[left]

            left += 1
        
        if right & 1:
            
            right -= 1

            m += tree[right]
        
        left //= 2
        right //= 2
    
    return m

def binary_search(tree,target,start,end):
    
    while start < end:
        
        mid = start + (end-start)//2

        k = query(tree,n,0,mid+1) #[0,mid]까지 구간합
        
        #구간합이 배치할 군인의 군번 target보다 크거나 같다면? mid에 배치가능
        if k >= target: 
            
            end = mid
        
        #target보다 작다면... 배치해야할 부대 번호는 mid + 1 이후에 있다
        else:
            
            start = mid + 1
    
    #index는 0번부터 세는데.. 부대번호는 1번부터 세야하니까..
    return end+1

n = int(stdin.readline())

A = list(map(int,stdin.readline().split()))

tree = [0]*(2*n)

create_segment(A,tree,n)

m = int(stdin.readline())

for _ in range(m):
    
    q = stdin.readline().rstrip()

    if q[0] == '1':
        
        _,b,c = map(int,q.split())

        update(tree,n,b-1,c)
    
    else:
        
        _,b = map(int,q.split())

        print(binary_search(tree,b,0,n))