문제

풀이

배열 길이가 2000 밖에 안 되길래 누적합만 구하고 직접 세도 되지 않을까 했는데... 효율성 테스트에서 걸러지고 말았다. 아무리 2000이라도 l, r, m 을 모두 반복문으로 구현하면 $O(n^3)$이니 시간초과가 발생한 모양이다. 시간복잡도를 낮추기 위해 m만 반복문으로 돌리고 l과 r은 투포인터 알고리즘을 사용하는 방법으로 다시 작성했다.

코드

def solution(cookie):
    n = len(cookie)
    answer = 0
    
    for m in range(n-1):
        l = m
        r = m + 1
        left = cookie[l]
        right = cookie[r]
        while l >= 0 and r < n:
            if left == right:
                answer = max(answer, left)
                l -= 1
                if l >= 0:
                    left += cookie[l]
            elif left < right:
                l -= 1
                if l >= 0:
                    left += cookie[l]
            else:
                r += 1
                if r < n:
                    right += cookie[r]
        
    return answer

과자 수가 무조건 양수이므로 투포인터 알고리즘을 적용할 수 있다. 이렇게 하면 시간복잡도가 $O(n^2)$이기 때문에 효율성 테스트까지 통과했다.

다만 포인터를 옮기는 과정과 합을 갱신하는 과정에서 조건 검사를 많이 해서 같은 조건을 반복해서 검사하고 있으며, 코드가 지저분하고 일관성 없어 보인다. 코드를 깔끔하게 바꾸기 위해서 다시 누적합으로 돌아왔다.

def solution(cookie):
    n = len(cookie)
    
    for i in range(1, n):
        cookie[i] += cookie[i-1]
    cookie = [0] + cookie
    
    answer = 0
    
    for m in range(1, n):
        l = m
        r = m + 1
        while l > 0 and r <= n:
            left = cookie[m] - cookie[l-1]
            right = cookie[r] - cookie[m]
            if left >= right:
                if left == right:
                    answer = max(answer, left)
                r += 1
            else:
                l -= 1
        
    return answer

누적합을 구해놓으면 합을 갱신하는 과정이 필요없기 때문에 좌우를 비교하는 과정과 포인터를 옮기는 과정을 깔끔하게 묶을 수 있다.

(오버헤드 때문인지 실행 시간은 조금 길어졌다.)