Problem p03070 (python)

Fixed submission:
def main():
	N = int(input())
	A = [0]*N
	MOD = 998244353
	summation = 0

	for i in range(N):
		A[i] = int(input())
		summation += A[i]

	# dp[i][j] represents sum of red is j at i position
	dp = [[0 for _ in range(summation+1)] for _ in range(N)]
	dp[0][0] = 2 # green or blue
	dp[0][A[0]] = 1 # red

	for i in range(1, N):
		for j in range(len(dp[i])):
			dp[i][j] += dp[i-1][j] * 2 % MOD # don't use A[i] for red, i.e. use green or blue
			dp[i][j] += dp[i-1][j-A[i]] % MOD if 0 <= j - A[i] <= len(dp[i]) else 0 # use i red

	# total
	ans = (3 ** N) % MOD
	ceil = summation // 2
	if summation % 2 != 0:
		ceil += 1
	ans -= sum([dp[N-1][r] for r in range(ceil, summation+1)])*3
	ans %= MOD

	if summation %2 == 0:
		# case r = g = S/2
		dp  =[[0 for _ in range(summation//2+1)] for _ in range(N)]
		dp[0][0] = 1 # only blue
		dp[0][A[0]] = 1
		for i in range(1, N):
			for j in range(len(dp[i])):
				dp[i][j] += dp[i-1][j] % MOD
				dp[i][j] += dp[i-1][j-A[i]] % MOD if 0 <= j-A[i] <= len(dp[i]) else 0
		ans += dp[N-1][summation//2]*3
		ans %= MOD

	print(ans)


if __name__ == '__main__':
    main()

Buggy submission (verdict: Runtime Error):
def main():
	N = int(input())
	A = [0]*N
	for i in range(N):
		A[i] = int(input())
 
	memo = dict()
	def create_triangles(r, g, b, num):
		ans = 0
		if num == N:
			if (abs(b-g)<r and r<b+g) or (abs(b-r)<g and g<b+r) or (abs(r-g)<b and b<r+g):
				return 1
			else:
				return 0
		if not (r+A[num],g,b) in memo:
			tmp = create_triangles(r+A[num], g, b, num+1)
			memo[(r+A[num],g,b)] = tmp
			memo[(r+A[num],b,g)] = tmp
			memo[(b,r+A[num],g)] = tmp
			memo[(b,g,r+A[num])] = tmp
			memo[(g,r+A[num],b)] = tmp
			memo[(g,b,r+A[num])] = tmp
		ans += memo[(r+A[num],g,b)]
		ans %= 998244353
 
		if not (r,g+A[num],b) in memo:
			tmp = create_triangles(r, g+A[num], b, num+1)
			memo[(r,g+A[num],b)] = tmp
			memo[(r,b,g+A[num])] = tmp
			memo[(b,r,g+A[num])] = tmp
			memo[(b,g+A[num],r)] = tmp
			memo[(g+A[num],r,b)] = tmp
			memo[(g+A[num],b,r)] = tmp
		ans += memo[(r,g+A[num],b)]
		ans %= 998244353
 
		if not (r,g,b+A[num]) in memo:
			tmp = create_triangles(r, g, b+A[num], num+1)
			memo[(r,g,b+A[num])] = tmp
			memo[(r,b+A[num],g)] = tmp
			memo[(b+A[num],r,g)] = tmp
			memo[(b+A[num],g,r)] = tmp
			memo[(g,r,b+A[num])] = tmp
			memo[(g,b+A[num],r)] = tmp
		ans += memo[(r,g,b+A[num])]
		ans %= 998244353
 
		return ans
 
	ans = create_triangles(0,0,0,0)
	print(ans)
 
 
if __name__ == '__main__':
    main()
def main():
	N = int(input())
	A = [0]*N
	for i in range(N):
		A[i] = int(input())

	memo = dict()
	def create_triangles(r, g, b, num):
		ans = 0
		if num == N:
			if (abs(b-g)<r and r<b+g) or (abs(b-r)<g and g<b+r) or (abs(r-g)<b and b<r+g):
				return 1
			else:
				return 0
		if not (r+A[num],g,b) in memo:
			tmp = create_triangles(r+A[num], g, b, num+1)
			memo[(r+A[num],g,b)] = tmp
			memo[(r+A[num],b,g)] = tmp
			memo[(b,r+A[num],g)] = tmp
			memo[(b,g,r+A[num])] = tmp
			memo[(g,r+A[num],b)] = tmp
			memo[(g,b,r+A[num])] = tmp
		ans += memo[(r+A[num],g,b)]
		ans %= 998244353

		if not (r,g+A[num],b) in memo:
			tmp = create_triangles(r, g+A[num], b, num+1)
			memo[(r,g+A[num],b)] = tmp
			memo[(r,b,g+A[num])] = tmp
			memo[(b,r,g+A[num])] = tmp
			memo[(b,g+A[num],r)] = tmp
			memo[(g+A[num],r,b)] = tmp
			memo[(g+A[num],b,r)] = tmp
		ans += memo[(r,g+A[num],b)]
		ans %= 998244353

		if not (r,g,b+A[num]) in memo:
			tmp = create_triangles(r, g, b+A[num], num+1)
			memo[(r,g,b+A[num])] = tmp
			memo[(r,b+A[num],g)] = tmp
			memo[(b+A[num],r,g)] = tmp
			memo[(b+A[num],g,r)] = tmp
			memo[(g,r,b+A[num])] = tmp
			memo[(g,b+A[num],r)] = tmp
		ans += memo[(r,g,b+A[num])]
		ans %= 998244353

		return ans

	ans = create_triangles(0,0,0,0)
	print(ans)


if __name__ == '__main__':
    main()