Implement Optimal Bucket Batching
Company: Amazon
Role: Machine Learning Engineer
Category: Coding & Algorithms
Difficulty: hard
Interview Round: Technical Screen
You are given an array `lengths` of `K` document lengths and an integer `G` representing the number of available GPUs. A batch is assigned to one GPU, and every document in the same batch is padded to the maximum document length in that batch. Reordering documents is allowed.
Choose up to `G` non-empty batches, with unused GPUs allowed, to minimize the total padding waste:
`sum over batches of (batch_size * max_length_in_batch - sum_of_lengths_in_batch)`
Return both the minimum padding waste and one optimal partition of the documents into batches. If `K = 0`, return zero waste and an empty partition. If `G >= K`, the optimal solution may place each document in its own batch, producing zero padding waste.
Implement an efficient algorithm and explain its time and space complexity.
Quick Answer: This question evaluates algorithm design and optimization skills, focusing on combinatorial partitioning and batching strategies, efficient utilization of parallel GPUs, and time/space complexity analysis relevant to ML data pipelines.
You are given an array `lengths` of document lengths and an integer `G` representing the number of available GPUs. You may reorder the documents, then split them into up to `G` non-empty batches. Every document in the same batch is padded to the maximum length in that batch, so a batch with lengths `[x1, x2, ..., xm]` creates padding waste `m * max(batch) - sum(batch)`. Return a tuple `(minimum_waste, partition)`, where `partition` is one optimal list of batches. Each batch should be returned as a list of lengths. Because reordering is allowed, returning batches in sorted order is acceptable. If multiple optimal partitions exist, return any one of them. If `lengths` is empty, return `(0, [])`.
Constraints
- 0 <= K == len(lengths) <= 5000
- 1 <= G <= 500 when K > 0
- 0 <= lengths[i] <= 10^9
Examples
Input: ([], 4)
Expected Output: (0, [])
Explanation: There are no documents, so there is no waste and no batches.
Input: ([7], 2)
Expected Output: (0, [[7]])
Explanation: A single document needs no padding.
Hints
- Because reordering is allowed, sort the lengths first. In an optimal solution, each batch can be taken as a contiguous segment of the sorted array.
- Use prefix sums to compute the waste of any sorted segment in O(1). Then define a DP over prefixes and number of batches; after algebraic rearrangement, each transition becomes a minimum over lines queried at `x = current_length`.