Skip to content

2551. Put Marbles in Bags 👍

  • Time: C++/Java: $O(\texttt{sort})$, Python: $O(n\log k)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
 public:
  long long putMarbles(vector<int>& weights, int k) {
    // To distribute marbles into k bags, there will be k - 1 cuts. If there's a
    // cut after weights[i], then weights[i] and weights[i + 1] will be added to
    // the cost. Also, no matter how we cut, weights[0] and weights[n - 1] will
    // be counted. So, the goal is to find the max/min k - 1 weights[i] +
    // weights[i + 1].
    vector<int> A;  // weights[i] + weights[i + 1]
    long long min = 0;
    long long max = 0;

    for (int i = 0; i + 1 < weights.size(); ++i)
      A.push_back(weights[i] + weights[i + 1]);

    ranges::sort(A);

    for (int i = 0; i < k - 1; ++i) {
      min += A[i];
      max += A[A.size() - 1 - i];
    }

    return max - min;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
  public long putMarbles(int[] weights, int k) {
    // To distribute marbles into k bags, there will be k - 1 cuts. If there's a
    // cut after weights[i], then weights[i] and weights[i + 1] will be added to
    // the cost. Also, no matter how we cut, weights[0] and weights[n - 1] will
    // be counted. So, the goal is to find the max/min k - 1 weights[i] +
    // weights[i + 1].
    int[] A = new int[weights.length - 1]; // weights[i] + weights[i + 1]
    long min = 0;
    long max = 0;

    for (int i = 0; i < A.length; ++i)
      A[i] = weights[i] + weights[i + 1];

    Arrays.sort(A);

    for (int i = 0; i < k - 1; ++i) {
      min += A[i];
      max += A[A.length - 1 - i];
    }

    return max - min;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
class Solution:
  def putMarbles(self, weights: List[int], k: int) -> int:
    # To distribute marbles into k bags, there will be k - 1 cuts. If there's a
    # cut after weights[i], then weights[i] and weights[i + 1] will be added to
    # the cost. Also, no matter how we cut, weights[0] and weights[n - 1] will
    # be counted. So, the goal is to find the max//min k - 1 weights[i] +
    # weights[i + 1].

    # weights[i] + weights[i + 1]
    A = [a + b for a, b in itertools.pairwise(weights)]
    return sum(heapq.nlargest(k - 1, A)) - sum(heapq.nsmallest(k - 1, A))