baekjoon_1126_two_towers python DP solution

two_towers.py

# raw_a = "35463 48571 46332 1 2 3 4 435 12 31 2 3 2 3 41 5 15 45433 34524 243 25345 2345 23534 52435 2435 2345 2345 24544 45445 4545 454 44544 545 454 34544 54 545 34 5324 5 6 78 5 4 32 6 735 72 5 5"

from collections import defaultdict 

def _solve(a, index, empty, tower1, tower2):
    """
    Dynamic Programming Solution
    """
    # base case
    if index == len(a):
        if tower1 == tower2 and tower1 != 0: 
            return tower1
        else:
            return -1
    
    # memoization
    if index in DP and abs(tower1 - tower2) in DP[index]:
        if DP[index][abs(tower1 - tower2)] == -1:
            return DP[index][abs(tower1 - tower2)]
        else:
            return DP[index][abs(tower1 - tower2)] + min(tower1, tower2)
    
    # DP cacluation
    choice1 = _solve(a, index + 1, empty + a[index], tower1, tower2)
    choice2 = _solve(a, index + 1, empty, tower1 + a[index], tower2)
    choice3 = _solve(a, index + 1, empty, tower1, tower2 + a[index])
    
    if choice1 != -1 : choice1 += min(tower1, tower2)
    if choice2 != -1 : choice2 += min(tower1, tower2)
    if choice3 != -1 : choice3 += min(tower1, tower2)
       
    DP[index][abs(tower1 - tower2)] = max(choice1, choice2, choice3)
#     print ("%d %d %d %d: %d" % (index, empty, tower1, tower2, DP[index][abs(tower1 - tower2)]))
    return DP[index][abs(tower1 - tower2)]
    

def main():  
    input()
    a = list(map(int, input().split()))
    print (_solve(a, 0, 0, 0, 0))

DP = defaultdict(dict)
main()