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> arr;  // weights[i] + weights[i + 1]
    long mn = 0;
    long mx = 0;

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

    ranges::sort(arr);

    for (int i = 0; i < k - 1; ++i) {
      mn += arr[i];
      mx += arr[arr.size() - 1 - i];
    }

    return mx - mn;
  }
};

 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[] arr = new int[weights.length - 1]; // weights[i] + weights[i + 1]
    long mn = 0;
    long mx = 0;

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

    Arrays.sort(arr);

    for (int i = 0; i < k - 1; ++i) {
      mn += arr[i];
      mx += arr[arr.length - 1 - i];
    }

    return mx - mn;
  }
}

 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]
    arr = [a + b for a, b in itertools.pairwise(weights)]
    return sum(heapq.nlargest(k - 1, arr)) - sum(heapq.nsmallest(k - 1, arr))