[백준 5676] 음주 코딩

문제 출처 : https://www.acmicpc.net/problem/5676

 

세그먼트 트리에 대한 이해가 필요한 문제 입니다. 기본적인 세그먼트 트리에 대해 아직 잘 모른다면 아래 링크를 통해 확인하시기 바랍니다.

https://wikidocs.net/209446

06. 세그먼트 트리

 

문제 이해하기

제목은 음주 코딩이지만 음주와는 전혀 관련 없는 문제 입니다. 주어진 수열이 있고 이 수열의 수들을 곱했을 때 양수가 될지, 음수가 될지, 0이 될지만 확인하면 됩니다.

다만 숫자가 너무 커질 것을 대비해서 직접 계산하지 않고 양수는 1, 음수는 -1, 0은 그대로 0으로 변경하면 숫자가 너무 커져서 생기는 문제를 해결할 수 있습니다.

그 외에는 크게 문제될 부분이 없습니다. 세그먼트 트리를 오랜만에 다시 풀어볼 때나, 이제 막 세그먼트 트리를 배웠다면 풀어보기 좋습니다.

코드 작성

그럼 직접 코드를 작성해 보겠습니다.

입력 받기

import sys
input = sys.stdin.readline

mii = lambda : map(int, input().split())

while True:
    try:
        N, K = mii()
        arr = list(mii())
        
        # 초기화 하기
        tree = [0] * (4 * N)

        for _ in range(K):
            cmd, a, b = input().split()
            a = int(a)
            b = int(b)
            
            # cmd를 통해 결과 값 찾기

    except:
        break

숫자가 많이 입력되기 때문에 sys.stdin.readline 을 통해 입력 받는 속도를 높여주었습니다. 다음으로 이 문제는 입력에 끝이 없습니다. 따라서 예외 처리를 통해 입력이 더 이상 들어오지 않으면 반복문을 빠져나오도록 하였습니다.

수열의 크기 N과 게임의 라운드 K를 입력 받습니다. 수열은 arr에 넣어주고, K개의 게임 명령어를 입력 받습니다. 세그먼트 트리를 만들기 위해서 Tree 리스트를 만들어 줍니다. 리스트의 정확한 크기로 계산할 필요가 없기 때문에 간단하게 N의 4배로 하면 됩니다.

명령 cmd는 C나 P이기 때문에 숫자로 받을 수 없습니다. 따라서 숫자로 변경하지 않고 입력 받은 것입니다. a와 b만 숫자로 변경해 줍니다.

세그먼트 트리 초기화

이제 세그먼트 트리를 만들기 위해 먼저 초기화 코드를 작성하겠습니다.

def change_tree_node(node, value):
    if 0 < value:
        tree[node] = 1
    elif value == 0:
        tree[node] = 0
    else:
        tree[node] = -1
    
    return tree[node]

def init(node, start, end):
    if start == end:
        tree[node] = arr[start]
        return tree[node]
    
    mid = (start + end) // 2
    i1 = init(2 * node, start, mid)
    i2 = init(2 * node + 1, mid + 1, end)
    return change_tree_node(node, i1 * i2)

세그먼트 트리 초기화 코드 입니다. 재귀 함수를 통해 트리에 값을 만들어 줍니다. 이 때 곱셈의 값이 너무 커질지 모르기 때문에 change_tree_node를 통해 결과값을 양수이면 1, 음수이면 -1, 0이면 0으로 변경해 주었습니다. 우리가 알고 싶은 것은 정확한 숫자가 아니기 때문에 이렇게 해도 상관 없습니다.

query 문 만들기

def query(node, start, end, left, right):
    if right < start or end < left:
        return 1
    
    if left <= start and end <= right:
        return change_tree_node(node, tree[node])
    
    mid = (start + end) // 2
    q1 = query(2 * node, start, mid, left, right)
    q2 = query(2 * node + 1, mid + 1, end, left, right)
    return q1 * q2

범위를 통해 곱셈의 결과를 가져오는 코드 입니다. 범위를 벗어난 경우 1을 리턴합니다. 0을 리턴하면 곱셈의 결과가 모드 0이 되기 때문에 1이 리턴되어야 합니다. 문제를 잘 보면 각 테스트 케이스에 곱셈 명령은 항상 한 번 이상있다. 라는 문구가 있습니다. 이 문구로 1을 리턴해도 된다는 것을 알 수 있습니다. 테스트 케이스에 곱셈이 있기 때문에 1로 리턴한 값이 결과값이 되는 것이 아니기 때문입니다.

update 문 만들기

def update(node, start, end, idx, value):
    if idx < start or end < idx:
        return
    
    if start == end:
        change_tree_node(node, value)
        return
    
    mid = (start + end) // 2
    update(2 * node, start, mid, idx, value)
    update(2 * node + 1, mid + 1, end, idx, value)
    rst = tree[2 * node] * tree[2 * node + 1]
    change_tree_node(node, rst)

update 입니다. start 와 end가 같아졌을 때 해당 인덱스에 있는 값을 value로 변경해 줍니다. 이 때도 숫자가 커지는 것을 놔둘 필요가 없기 때문에 위에서 사용한 change_tree_node 함수로 1, -1, 0으로 변경해 줍니다. update 후에는 변경된 값으로 모든 노드들의 결과를 수정해 줍니다.

결과 출력 하기

while 문 전체 입니다. 정확히 어느 위치에 코드가 들어가야 할지 모를 수 있기 때문에 전체 코드로 작성하였습니다.

while True:
    try:
        N, K = mii()
        arr = list(mii())
        tree = [0] * (4 * N)

        init(1, 0, N - 1)

        ans = ''
        for _ in range(K):
            cmd, a, b = input().split()
            a = int(a)
            b = int(b)

            if cmd == 'C':
                update(1, 0, N - 1, a - 1, b)
            else:
                rst = query(1, 0, N - 1, a - 1, b - 1)
                if 0 < rst:
                    ans += '+'
                elif rst == 0:
                    ans += '0'
                else:
                    ans += '-'

        print(ans) 

    except:
        break

cmd값이 C이면 update를, P이면 query 문을 수행 합니다. 이 때 a, b값은 인덱스가 된다면 1을 빼주고, update의 value 값이 될 때는 그대로 사용 됩니다. 그래서 a, b 값에 -1을 해준 부분이 있는 것입니다.

query 함수를 통해 나온 결과값이 0보다 크면 +. 음수면 -, 0인 경우 0으로 변경하여 ans에 이어붙여 주었습니다. 모든 게임 라운드가 끝나면 그 결과값을 출력해 주면 됩니다.

전체 코드

전체 코드를 확인해 보겠습니다.

import sys
input = sys.stdin.readline

mii = lambda : map(int, input().split())

def change_tree_node(node, value):
    if 0 < value:
        tree[node] = 1
    elif value == 0:
        tree[node] = 0
    else:
        tree[node] = -1
    
    return tree[node]

def init(node, start, end):
    if start == end:
        tree[node] = arr[start]
        return tree[node]
    
    mid = (start + end) // 2
    i1 = init(2 * node, start, mid)
    i2 = init(2 * node + 1, mid + 1, end)
    return change_tree_node(node, i1 * i2)

def query(node, start, end, left, right):
    if right < start or end < left:
        return 1
    
    if left <= start and end <= right:
        return change_tree_node(node, tree[node])
    
    mid = (start + end) // 2
    q1 = query(2 * node, start, mid, left, right)
    q2 = query(2 * node + 1, mid + 1, end, left, right)
    return q1 * q2

def update(node, start, end, idx, value):
    if idx < start or end < idx:
        return
    
    if start == end:
        change_tree_node(node, value)
        return
    
    mid = (start + end) // 2
    update(2 * node, start, mid, idx, value)
    update(2 * node + 1, mid + 1, end, idx, value)
    rst = tree[2 * node] * tree[2 * node + 1]
    change_tree_node(node, rst)

while True:
    try:
        N, K = mii()
        arr = list(mii())
        tree = [0] * (4 * N)

        init(1, 0, N - 1)

        ans = ''
        for _ in range(K):
            cmd, a, b = input().split()
            a = int(a)
            b = int(b)

            if cmd == 'C':
                update(1, 0, N - 1, a - 1, b)
            else:
                rst = query(1, 0, N - 1, a - 1, b - 1)
                if 0 < rst:
                    ans += '+'
                elif rst == 0:
                    ans += '0'
                else:
                    ans += '-'

        print(ans) 

    except:
        break