[BOJ 27210] 신을 모시는 사당 - Python 풀이

링크

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

문제

신을 모시는 사당에는 신을 조각한 돌상 N개가 일렬로 놓여 있다. 각 돌상은 왼쪽 또는 오른쪽을 바라보고 서있다. 창영이는 연속한 몇 개의 돌상에 금칠을 하여 궁극의 깨달음을 얻고자 한다.

궁극의 깨달음을 얻기 위해서는 가능한 한 많은 금색 돌상들이 같은 방향을 바라보아야 한다. 방향이 다른 돌상은 깨달음에 치명적이다. 깨달음의 양은 아래와 같이 정의된다.

| (왼쪽을 바라보는 금색 돌상의 개수) - (오른쪽을 바라보는 금색 돌상의 개수) |

창영이는 궁극의 깨달음을 얻을 수 있을까?

입력

첫째 줄에 돌상의 개수 N이 주어진다.

둘째 줄에 돌상이 나열된 순서대로, 각 돌상이 바라보고 있는 방향이 주어진다. 입력의 편의상 왼쪽은 1, 오른쪽은 2라고 하자.

출력

최대한 많은 깨달음을 얻기 위해 금을 칠하였을 때, 얻을 수 있는 깨달음의 양을 출력한다.

제한

• 1 ≤ N ≤ 100,000

코드

import sys

N = int(sys.stdin.readline())
arr = list(map(int, sys.stdin.readline().split()))
prefix_1 = [0] * N
prefix_2 = [0] * N

prefix_1[0] = 1 if arr[0] == 1 else 0
for i in range(1, N):
    if arr[i] == 1: 
        prefix_1[i] = prefix_1[i-1]+1
    else: 
        prefix_1[i] = prefix_1[i-1]-1 if prefix_1[i-1] > 0 else 0

prefix_2[0] = 1 if arr[0] == 2 else 0
for i in range(1, N):
    if arr[i] == 2: 
        prefix_2[i] = prefix_2[i-1]+1
    else: 
        prefix_2[i] = prefix_2[i-1]-1 if prefix_2[i-1] > 0 else 0

max_val = 0

print(max(max(prefix_1), max(prefix_2)))

풀이

대회 당시에는 못풀었다ㅠ 그러나 외양간 고치는 마음으로 바로 풀어보았다.

먼저 prefix_1과 prefix_2라는 배열을 두어 각각 왼쪽을 보는 불상의 구간합, 오른쪽을 보는 불상의 구간합을 구하고자 했다.

prefix_1은 1의 구간합을 구하는 배열이다. arr[i]가 1이 나올 경우 이전 구간합에다 1을 더한 값을 저장했고, 1이 나오지 않았을 경우에는 이전 구간합에서 -1을 뺀 값을 저장했다. 이때 구간합이 음수가 되면 안되므로 0보다 큰 경우만 -1을 빼주었고 0보다 작은 경우에는 0이 되도록 했다.

prefix_2는 2의 구간합을 구하는 배열이며 로직은 prefix_1과 같이 구했다.

prefix_1, prefix_2를 모두 구한 후에는 각 배열의 최대값을 구한 후 둘 중 큰 값을 선택하여 출력했다.