[Python][백준] 1321번 군인

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

공부 날짜 : 2023.02.16
정답 참조 여부 : X

군대에서 각 부대에 배정할 인원이 계속해서 바뀌며 각 군인은 부대에 배정인원에 따라 순서대로 배정된다. 군인의 번호가 주어졌을 때, 군인이 배정되는 부대를 구하는 문제이다.

---

단순히 구현만 하면 각 부대의 배정 인원을 구하고, 1번 부대에서부터 배정되는 인원을 순차로 빼주면 될 것이다.

하지만 그렇게 하면 당연히 시간초과가 난다.

그래서 떠올린 방법은 2진 탐색처럼 중앙에서부터 누적합을 구하는 방법이였는데,
이게 세그먼트 트리 라고 한다.

구현한 방식은

1 4 3 5 2
5 8 2
13 2

이런 식으로 1,2번 부대 배정 병사수, 2,3번 부대 배정 병사수, 5번 부대 배정 병사 수로 합쳐주고

다시 1,2,3,4번 부대 배정 병사수, 5번 부대 배정 병사수로 나눴다.

트리의 형태를 의도한건데 그림을 못그리고 데이터로 표현해서 잘 표현이 안된다.

다시 설명하자면 1,2,3,4번 부대 배정되는 병사수의 자식노드로 1,2번, 3,4번
1,2번 부대 배정되는 병사수의 자식노드로 1번, 2번

이런식으로 구현했다.

그러면 군인의 번호로 부대까지 찾아가는데 O(logn)으로 탐색 가능해 진다.

소스코드

import sys
input = sys.stdin.readline
#####################################################
n = int(input())

unit_tree = [list(map(int, input().split()))]


# 입력받은 데이터를 트리형태로 저장하기
a = n//2
while True:
    if a == 0:
        break

    new_unit = [0 for _ in range(a+1)]
    for i in range(len(unit_tree[-1])):
        new_unit[i//2] += unit_tree[-1][i]

    unit_tree.append(new_unit.copy())

    a //= 2


########################################################
m = int(input())

for _ in range(m):
    order, *detail = map(int, input().split())

    if order == 1:
        detail[0] -= 1
        i = 0
        while True:
            if i == len(unit_tree):
                break
            unit_tree[i][detail[0]] += detail[1]
            detail[0] = (detail[0]) // 2
            i += 1

        continue


    # 여기서 부턴 출력만 해보자

    i = len(unit_tree) - 1
    j = 0

    while True:
        if i == -1:
            print(j//2+1)
            break

        if unit_tree[i][j] >= detail[0]:
            i -= 1
            j *= 2
            continue

        if unit_tree[i][j] < detail[0]:
            detail[0] -= unit_tree[i][j]
            j += 1