[GO] Precautions when solving DP problems with Python

Overview

When I was solving the DP (Dynamic Programming) problem of AtCoder, I encountered a phenomenon that the code that would be AC in other languages would be TLE if it was Python / PyPy, so I will leave it as a memorandum.

About DP problem

As explained in detail in this article, there are several ways to solve the DP problem, but roughly speaking, the for statement is turned bottom-up. There are two methods, one is to build a DP table and the other is memoization recursion, which calls recursive functions top-down. In any case, it is the same in that the processing speed is increased by retaining the calculated result, but in the case of memoization recursion, there is an overhead of function call, so it is within the time limit in some languages. In some cases, the processing is not completed and it becomes TLE.

example

DP Summary Contest B Problem: Frog2

Solution 1: Bottom-up DP by for statement

N, K = map(int, input().split())
h = list(map(int, input().split()))
dp = [10**9] * N

for i in range(N):
    if i == 0:
        dp[0] = 0
    elif i == 1:
        dp[1] = abs(h[0]-h[1])
    else:
        for j in range(1, K+1):
            if j > i:
                break
            dp[i] = min(dp[i], dp[i - j] + abs(h[i - j] - h[i]))

print(dp[N-1])

result

It becomes TLE when submitted in Python and AC when submitted in PyPy ). You can also get AC by writing similar code in a fast language such as Go.

Solution 2: Memoization recursion (number of function calls: large)

When written in Python

import sys
sys.setrecursionlimit(10**6)

N, K = map(int, input().split())
h = list(map(int, input().split()))
INF = 10**9
dp = [INF] * N

def solve(n):
    if dp[n] != INF:
        return dp[n]

    if n == 0:
        dp[n] = 0
    elif n == 1:
        dp[n] = abs(h[0] - h[1])
    else:
        for i in range(1, K + 1):
            if n - i < 0:
                break
            cost = solve(n - i) + abs(h[n - i] - h[n])
            dp[n] = min(dp[n], cost)

    return dp[n]

print(solve(N-1))

When written in Go

package main

import (
	"fmt"
	"math"
)

var (
	N, K, INF int
	h, dp []int
)

func solve(n int) int {
	if dp[n] != INF {
		return dp[n]
	}

	if n == 0 {
		dp[n] = 0
	} else if n == 1 {
		dp[n] = int(math.Abs(float64(h[n] - h[n-1])))
	} else {
		for i := 1; i <= K; i++ {
			if n - i < 0 {
				break
			}
			var cost =  solve(n - i) + int(math.Abs(float64(h[n-i]-h[n])))
			dp[n] = int(math.Min(float64(dp[n]), float64(cost)))
		}
	}

	return dp[n]
}

func main() {
	fmt.Scan(&N, &K)
	h = make([]int, N)
	dp = make([]int, N)
	INF = int(math.Pow(10, 9))

	for i := 0; i < N; i++ {
		fmt.Scan(&h[i])
	}
	for i := 0; i < N; i++ {
		dp[i] = INF
	}

	fmt.Println(solve(N - 1))
}

result

It is a so-called ordinary memoization recursion, but it becomes TLE when submitted with Python or PyPy, and [Write and submit in a high-speed language such as Go]. Then it becomes AC](https://atcoder.jp/contests/dp/submissions/14452835). In order to take AC with Python / PyPy, another device is required.

Solution 3: Memoization recursion (number of function calls: small)

import sys
sys.setrecursionlimit(10**6)

N, K = map(int, input().split())
h = list(map(int, input().split()))
INF = 10**9
dp = [INF] * N

def solve(n):
    if dp[n] != INF:
        return dp[n]

    if n == 0:
        dp[n] = 0
    elif n == 1:
        dp[n] = abs(h[0] - h[1])
    else:
        for i in range(1, K+1):
            if n - i < 0:
                break
            if dp[n - i] != INF:
                #Dp to reduce the number of function calls[n-i]Use it if is already calculated
                cost = dp[n - i] + abs(h[n - i] - h[n])
            else:
                cost = solve(n - i) + abs(h[n - i] - h[n])
            dp[n] = min(dp[n], cost)

    return dp[n]

print(solve(N-1))

result

The code is almost the same as Solution 2, but in this case, it checks whether dp [n-i] has been calculated before calling solve (n-i), and if it is calculated, it is used. By doing this, you can reduce the overhead of function call and make it AC with PyPy. However, even with this ingenuity, it will be TLE when submitted in Python.

Summary

The results for each solution and language are as follows. In other words, if you submit it in Python, if you devise TLE and PyPy, you can take AC without being particularly conscious of high-speed languages such as AC and Go.

So it may be a good idea to submit DP issues in a fast language such as C ++ or Go. If you really want to solve it with Python, you should be aware of the following.

--If possible, do bottom-up DP with a for statement instead of memoization recursion. --When performing memoization recursion, reduce the number of function calls as much as possible (check the DP table before calling the function) --Submit in PyPy instead of Python

If you have any other mistakes such as points you noticed or recognition, it would be helpful if you could point them out.