LeetCode-in-Java

2551. Put Marbles in Bags

Hard

You have k bags. You are given a 0-indexed integer array weights where weights[i] is the weight of the ith marble. You are also given the integer k.

Divide the marbles into the k bags according to the following rules:

The score after distributing the marbles is the sum of the costs of all the k bags.

Return the difference between the maximum and minimum scores among marble distributions.

Example 1:

Input: weights = [1,3,5,1], k = 2

Output: 4

Explanation:

The distribution [1],[3,5,1] results in the minimal score of (1+1) + (3+1) = 6.

The distribution [1,3],[5,1], results in the maximal score of (1+3) + (5+1) = 10.

Thus, we return their difference 10 - 6 = 4.

Example 2:

Input: weights = [1, 3], k = 2

Output: 0

Explanation: The only distribution possible is [1],[3]. Since both the maximal and minimal score are the same, we return 0.

Constraints:

Solution

import java.util.Arrays;

public class Solution {
    public long putMarbles(int[] weights, int k) {
        long minAns = weights[0] + (long) weights[weights.length - 1];
        long maxAns = weights[0] + (long) weights[weights.length - 1];
        long[] arr = new long[weights.length - 1];
        for (int i = 1; i < weights.length; i++) {
            arr[i - 1] = weights[i] + (long) weights[i - 1];
        }
        Arrays.sort(arr);
        for (int i = 0; i < k - 1; i++) {
            minAns = minAns + arr[i];
        }
        for (int i = arr.length - 1; i > arr.length - k; i--) {
            maxAns = maxAns + arr[i];
        }
        return maxAns - minAns;
    }
}