[ BOJ ] 27651 : 벌레컷 ( GOLD 3 ) / Python

문제

크기 N의 1차원 양의 정수 배열로 이루어진 자벌레가 있다. 자벌레는 곤충이기 때문에 머리, 가슴, 배로 부위를 구분할 수 있다.

각 부위는 배열상에서 연속하는 구간으로 나타낼 수 있으며 배열상에서 머리는 왼쪽에, 가슴은 가운데에, 배는 오른쪽에 존재한다. 각 부위의 크기는 배열상에서 해당하는 구간의 값의 합으로 정의된다.

무지는 이 자벌레가 가슴이 배보다 크고 배가 머리보다 크다는 사실은 알고 있지만 어느 지점에서 머리 가슴 또한 가슴 배가 구분되는지 알지 못한다. 무지를 도와 구분될 수 있는 경우의 수를 구해주자.

엄밀하게는 다음 조건에 맞는 X,Y 쌍의 개수를 구해주자.

배열에 i번째 원소의 값을 A_i라 할 때

입력

첫 줄에 정수 N이 주어진다. (3 ≤ N ≤ 1 000 000)

두 번째 줄에 배열의 값 A_1, A_2, ... , A_N이 공백으로 구분되어 주어진다. (1 ≤ A_i ≤ 100 000)

출력

자벌레의 머리 가슴 배가 구분될 수 있는 경우의 수를 구해 출력하자.

풀이 과정

머리[0..X] < 배 [Y+1..N-1] < 가슴 [X+1..Y]

 

먼저, Y를 정해놓은 뒤 누적합 배열에서 X를 이분 탐색으로 찾겠다는 풀이를 떠올렸다. 

머리 / 가슴 / 배는 각각 해당 구간의 합이고, Y가 특정 값일 때 머리[0..X] < 배 [Y+1..N-1] < 가슴 [X+1..Y] 관계가 성립하는 X의 개수 를 브루트포스로 찾기에는 시간이 오래 걸리므로, 누적합 이분탐색을 사용했다.

 

X가 커질수록 머리 구간의 크기가 증가하고 가슴 구간의 크기가 감소한다.

X가 작아질수록 머리 구간의 크기가 감소하고 가슴 구간의 크기가 증가한다.

 

배 구간의 크기는 정해져 있으므로, X를 크거나 작게 만들어 가능한 경우를 탐색한다.

 

탐색을 하면서 만약 가능한 X를 이분 탐색의 mid 를 통해 찾았다면, X가 작아져도 관계는 무조건 성립하기 때문에 ( 머리 구간의 크기가 감소하고 가슴 구간의 크기가 증가하므로 머리 < 배 < 가슴 관계를 계속 유지함. ) mid 를 결과값에 더해준다. 이후 탐색을 더 진행한다면 해당 mid 값은 증가시켰으므로 이후에 결과값에 중복으로 더하지 않는다. ( already 변수를 사용해 처리했다. )

 

import sys
input = sys.stdin.readline

# 머리 / 가슴 / 배
# 머리 < 배 < 가슴

# 1. y 하나 정해놓기
# 2. mid까지 누적 합 < y까지의 누적 합 < y까지의 누적 합 - mid까지 누적 합 로 x를 찾으면서, 머리가 배보다 안 커질때까지의 개수 구하기

n = int(input().rstrip())
arr = list(map(int, input().rstrip().split()))

sums = [0]
for a in arr: sums.append(sums[-1] + a)
del sums[0]

result = 0
for y in range(1, n - 1):
    stomSum = sums[-1] - sums[y] # y까지의 누적합

    start = 0
    end = y - 1
    if start == end:
        if sums[0] < stomSum < sums[y] - sums[0]: result += 1
    else:
        cnt = 0
        already = 0
        while start <= end:
            mid = (start + end) // 2
            if sums[mid] < stomSum < sums[y] - sums[mid]:
                result += mid + 1 - already # mid + 1 부터 계속 더 탐색하기
                already = mid + 1
                start = mid + 1
            else:
                end = mid - 1

print(result)