Given n stacks of ints, find the maximum sum of k numbers by popping out any stacks


  • 0
    R

    Assuming you have n stacks, find the maximum sum of k numbers. You can do the standard stack operations.


  • 0
    R
    struct Node{
      int val;
      int index;
      bool operator<(const Node& n) const{
        return val > n.val;
      }
    };
    
    int MaxK(vector<stack<int>> v, int k){
      priority_queue<Node> pq;
      for(int i=0;i<n;i++){
        auto e = v[i];
        if(e.empty()) continue;
        pq.push({e.top(), i});
      }
      int ans= 0;
      for(int i=0;i<k;i++){
        auto cur = pq.top();
        ans+=cur.fi;pq.pop();
        v[cur.se].pop();
        if(v[cur.se].empty()) continue;
        pq.push(v[cur.se].top());
      }
      return ans;
    }
    

    I think is pq problem, why is tagged as dp?


  • 0
    R

    Is the solution correct? Say the first stack is [100,1,2], the second stack is [7,8,10], if the k = 3, your solution will result in 25 (7 + 8 + 10) if I understand correctly, but the true answer is 103 (100 + 1 + 2)


  • 0
    A

    Sort and merge all n stack into one stack (which similar to sort one stacks use an additional stack), and then pop k numbers ans sum them up?
    How to apply DP?


  • 0
    S

    use minheap of size k....pop ele.s from stacks and insert in heap....O(n.logk)


  • 0
    M

    The question was not stated well. If you pop an element off a stack you must include that element as part of the k numbers that comprise the final sum.

    royzxq's example shows that.

    The PQ answers don't work.


  • 0
    A

    @mwistrom I think you are right.

    My DP solution in Python for this interpretation with O(n*k^2) complexity :

    # max(a + total, b) that deals with None values
    def smart_max(total, a, b):
        if a is None:
            return b
        if b is None:
            return total + a
    
        return max(total + a, b)
    
    
    def find_solutions(stacks, k):
        n = len(stacks)
        if n == 0:
            return None
    
        S = [[None for _ in range(k + 1)] for __ in range(n)]
        S[0][0] = 0
        for i in range(1, min(k, len(stacks[0])) + 1):
            S[0][i] = S[0][i - 1] + stacks[0][-i]
    
        for p in range(1, n):
            S[p][0] = 0
            total = 0
            for i in range(0, min(k, len(stacks[p])) + 1):
                for j in range(i, k + 1):
                    S[p][j] = smart_max(total, S[p - 1][j - i], S[p][j])
                if i < len(stacks[p]):
                    total += stacks[p][-i - 1]
    
        return S[n - 1][k]
    

    And the test cases :

    def test_solution(stacks, k, solution):
        actual = find_solutions(stacks, k)
        assert actual == solution
    
    
    stacks = [
        [9, -3],
        [-2, 4, 6],
        [],
        [1, 2],
        [3]
    ]
    
    test_solution([], 1, None)
    
    test_solution(stacks[:1], 0, 0)
    test_solution(stacks[:2], 0, 0)
    test_solution(stacks[:3], 0, 0)
    test_solution(stacks[:4], 0, 0)
    test_solution(stacks[:5], 0, 0)
    
    test_solution(stacks[:1], 1, -3)
    test_solution(stacks[:2], 1, 6)
    test_solution(stacks[:3], 1, 6)
    test_solution(stacks[:4], 1, 6)
    test_solution(stacks[:5], 1, 6)
    
    test_solution(stacks[:1], 2, 6)
    test_solution(stacks[:2], 2, 10)
    test_solution(stacks[:3], 2, 10)
    test_solution(stacks[:4], 2, 10)
    test_solution(stacks[:5], 2, 10)
    
    test_solution(stacks[:1], 3, None)
    test_solution(stacks[:2], 3, 12)
    test_solution(stacks[:3], 3, 12)
    test_solution(stacks[:4], 3, 12)
    test_solution(stacks[:5], 3, 13)
    
    stacks = [
        [100, 1, 2],
        [7, 8, 10]
    ]
    
    test_solution(stacks[:1], 1, 2)
    test_solution(stacks[:2], 1, 10)
    
    test_solution(stacks[:1], 2, 3)
    test_solution(stacks[:2], 2, 18)
    
    test_solution(stacks[:1], 3, 103)
    test_solution(stacks[:2], 3, 103)
    

    I feel like we could do better that O(n*k^2), but for a small enough k this is good.

    EDIT : little explanation, S[p][i] is the solution when using only stacks up to stack p and i numbers (instead of k)
    Note that we only use S[p-1] at each iteration, so only keeping the last line in memory would reduce space complexity, I did not do it cause I think it is easier to understand this way.


  • 0
    M

    O(n * k^2) is the best complexity I could come up with too. I'm thinking that is the best there is. O(n * k) is required at a minimum to check out all the possible values that could be selected. Can't really see a way to get rid of that third inner loop. It is quite a nice little problem though.


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.