Optimal Stock Selection Algorithm (Python)

Question

Given a total amount of available capital, the values of stocks at the current date and the values of the same stocks at a later date, I need to find a way to calculate the optimal return if I can buy at most one share of any stock.

For example:

Given

input:
capital = 10
cur_vals = [3,4,5,6]
fut_vals = [6,7,5,3]

output:
6

So far I've come up with the following:

def optimal_select(spend, cur_value, fut_value):
    
    #create list of (returns,current_vals)
    
    diffs = sorted([(fut_value[i] - cur_value[i],cur_value[i]) for i in range(len(cur_value))], reverse = True)
    
    values = []
    
    for k in range(len(diffs)):
        
        #initial balance:
        balance = spend - diffs[k][1]      
        
        #initial returns
        returns = diffs[k][0]
        
        #attempt to add each additional stock by highest return
        
        for j in range(k+1,len(diffs)):
            
            if diffs[j][1] <= balance and diffs[j][1] > 0:
                
                balance += -diffs[j][1]
                returns += diffs[j][0]
        
        values.append(returns)
    
    return max(values)

Although, I know this won't give me the correct answer all the time because it doesn't account for the case where it makes sense to forgo the highest return selection at a given iteration for multiple less costly stocks that sum to higher returns. For example:

optimal_select(8, [4,4,2,2],[10,9,5,5])

output: 11

Where the correct output in this case should be 12.

Is there a particular algorithm or area of optimization that I can reference that might help me better structure my approach?

Answer 1

Thanks to the comment, I was able to work through the correct solution with the 0-1 knapsack problem and the following code:

def optimal_selection(w, cost, returns):

    m = [[None for i in range(len(cost)+1)] for j in range(w + 1)]

    for j in range(w+1):
        m[j][0] = 0

    for i in range(1,len(cost)+1):

        for j in range(w+1):

            if cost[i-1] > j:

                m[j][i] = m[j][i-1]

            else:

                m[j][i] = max(m[j][i-1], m[j-cost[i-1]][i-1] + returns[i-1])

            
    return max([i for sub in m for i in sub])

```