#P4625.第2题-大模型训练显存优化算法

第2题-大模型训练显存优化算法 - problem_ide - CodeFun2000

动态规划

import sys
import numpy
def solve():
    data = sys.stdin.read().split()
    if not data:
        return
    
    m = int(data[0])
    n = int(data[1])

    space = [int(x) for x in data[2 : n+2]]
    swap = [int(x) for x in data[n+2 : 2*n+2]]
    price = [int(x) for x in data[2*n+2 : 3*n+2]]

    cost = [min(swap[i], price[i]) for i in range(n)]

    # 3. 初始化 DP 数组
    INF = float('inf')
    dp = [INF] * (m + 1)
    dp[0] = 0

    for i in range(n):

        for j in range(m,-1,-1):
            if dp[j] != INF:
                nj = j + space[i]
                if j + space[i]>= m:
                    nj = m
                
                if dp[j]+cost[i]<dp[nj]:
                    dp[nj] = dp[j]+cost[i]
    
    if dp[m]==INF:
        print('error')
    else:
        print(dp[m])

if __name__=='__main__':
    solve()

                

递归(超时)

import sys
import numpy


    
def solve():
    data = sys.stdin.read().split()
    if not data:
        return
    
    m = int(data[0])
    n = int(data[1])

    space = [int(x) for x in data[2 : n+2]]
    swap = [int(x) for x in data[n+2 : 2*n+2]]
    price = [int(x) for x in data[2*n+2 : 3*n+2]]

    cost = [min(swap[i], price[i]) for i in range(n)]

    def dfs(step,now_weight,now_cost):
        if step==n:
            global ans
            if now_weight>=m:
                if ans == -1:
                    ans = now_cost
                else:
                    ans = min(ans,now_cost)
            
            return
        
        dfs(step+1,now_weight,now_cost)
        dfs(step+1,now_weight+space[step],now_cost+cost[step])
    
    dfs(0,0,0)
    if ans==-1:
        print('error')
    else:
        print(ans)
        



if __name__=='__main__':
    ans = -1
    solve()