Quick Overview

Build an ordered vector data pipeline with batch-local right padding, empty-vector handling, and optional dropping of the final incomplete batch.

Batch and Pad Variable-Length Vectors with Optional Final-Batch Dropping

Company: Siemens

Role: Machine Learning Engineer

Category: Coding & Algorithms

Difficulty: medium

Interview Round: Technical Screen

Implement a deterministic data pipeline that groups variable-length integer vectors into batches, pads each retained batch, and optionally drops the final incomplete batch. The exercise models three responsibilities: `create_data_iterator` groups vectors in input order, `pad_vector_collate` right-pads one batch to its longest vector, and `run_data_pipeline` returns the collected batches. Only the final function is the console entry point; internal helper organization is up to you. ### Function Signature `run_data_pipeline(vectors: list[list[int]], batch_size: int, pad_value: int, drop_last: bool) -> list[list[list[int]]]` ### Rules - Read vectors in their original order without shuffling or mutation. - Form consecutive groups of `batch_size` vectors. - If the final group contains fewer vectors and `drop_last` is true, discard that entire group. Otherwise retain it. - For each retained group, let `L` be the greatest original vector length in that group. Append `pad_value` to each vector until its length is `L`. - Padding is batch-local, not based on the longest vector in the whole dataset. - Empty vectors are valid. A retained batch containing only empty vectors contains the same number of empty vectors after padding. These padding direction, pad-value input, ordering, and incomplete-batch rules are explicit conventions for this exercise. ### Output Return the padded batches in input order. Return `[]` if no batch is retained. ### Constraints - `0 <= len(vectors) <= 10000`. - `1 <= batch_size <= 1000`. - `0 <= len(vector) <= 1000` for each vector. - Vector elements and `pad_value` are integers in `[-1000000, 1000000]`. - The total number of input elements and the total number of elements after padding are each at most `1000000`. - All input arguments are valid. ### Examples Input: `vectors = [[1,2],[3],[],[4,5,6],[7]], batch_size = 2, pad_value = 0, drop_last = false` Output: `[[[1,2],[3,0]],[[0,0,0],[4,5,6]],[[7]]]` With the same arguments except `drop_last = true`, output: `[[[1,2],[3,0]],[[0,0,0],[4,5,6]]]` Input: `vectors = [[],[]], batch_size = 2, pad_value = -1, drop_last = true` Output: `[[[],[]]]` Input: `vectors = [], batch_size = 3, pad_value = 0, drop_last = false` Output: `[]`

Overview: Build an ordered vector data pipeline with batch-local right padding, empty-vector handling, and optional dropping of the final incomplete batch.

Read the full Siemens Machine Learning Engineer interview experience this question came from

Implement a deterministic data pipeline that groups variable-length integer vectors into batches, pads each retained batch, and optionally drops the final incomplete batch. The exercise models three responsibilities: grouping the vectors in input order, right-padding one batch to its longest vector, and returning the collected batches. Only `run_data_pipeline` is the console entry point; how you organize any internal helpers is up to you. ### Rules 1. Read the vectors in their original order, without shuffling or mutating them. 2. Form consecutive groups of `batch_size` vectors: the first `batch_size` vectors, then the next `batch_size`, and so on. 3. If the final group contains fewer than `batch_size` vectors and `drop_last` is true, discard that entire group. Otherwise retain it. 4. For each retained group, let `L` be the greatest original vector length in that group. Append `pad_value` to each vector of the group until its length is exactly `L`. Padding is appended on the right, never prepended. 5. Padding is batch-local: `L` is the longest vector of that group only, not the longest vector of the whole dataset. 6. Empty vectors are valid. A retained batch containing only empty vectors has `L = 0`, so it contains the same number of empty vectors after padding. These padding-direction, pad-value, ordering, and incomplete-batch rules are explicit conventions of this exercise. ### Output Return the padded batches as a list of batches, in input order. Return `[]` if no batch is retained. Every value that crosses the boundary is an element or `pad_value` in `[-1000000, 1000000]`; no value can exceed 2^31 - 1, so Java `int` and C++ `int` are sufficient (no `long`/`long long` is needed). ### Examples Example 1: `vectors = [[1, 2], [3], [], [4, 5, 6], [7]]`, `batch_size = 2`, `pad_value = 0`, `drop_last = False`. Output: `[[[1, 2], [3, 0]], [[0, 0, 0], [4, 5, 6]], [[7]]]`. The groups are `[[1,2],[3]]` with `L = 2`, `[[],[4,5,6]]` with `L = 3`, and the final short group `[[7]]` with `L = 1`, which is retained because `drop_last` is false. With the same arguments except `drop_last = True`, the output is `[[[1, 2], [3, 0]], [[0, 0, 0], [4, 5, 6]]]` because the final group holds fewer than `batch_size` vectors. Example 2: `vectors = [[], []]`, `batch_size = 2`, `pad_value = -1`, `drop_last = True`. Output: `[[[], []]]`. The single group is full, its greatest length is 0, so no padding is appended.

Constraints

  • 0 <= len(vectors) <= 10000.
  • 1 <= batch_size <= 1000.
  • 0 <= len(vector) <= 1000 for each vector.
  • Vector elements and pad_value are integers in [-1000000, 1000000].
  • The total number of input elements and the total number of elements after padding are each at most 1000000.
  • drop_last is a boolean.
  • All input arguments are valid.
  • No value crossing the interface can exceed 2^31 - 1, so 32-bit integers suffice in every language.

Examples

Input: ([], 3, 0, False)

Expected Output: []

Explanation: Minimum valid input: no vectors means no group is formed, so the result is empty even when drop_last is false.

Input: ([], 3, 0, True)

Expected Output: []

Explanation: Empty input with drop_last true also yields no batches; there is nothing to drop.

Hints

  1. Only one group can ever be incomplete. Which one is it, and what does that tell you about when the drop_last test actually has to run?
  2. The pad target L is defined by the original lengths inside a single group, so decide L before you write any padding and make sure a group cannot see padding produced for an earlier group.
  3. Work through the cases the rules single out: an empty vectors list, a batch whose vectors are all empty, and an empty vector sitting next to a long one in the same batch.

Loading coding console...

Show the approach

Approach

Algorithm: walk the input with a cursor that advances by batch_size, so the slice [start, min(start + batch_size, n)) is the next consecutive group in input order. Because the cursor advances by a fixed stride from index 0, only the last group can contain fewer than batch_size vectors; when that happens and drop_last is true, the group is skipped entirely and nothing derived from it is emitted. For a retained group, a single pass computes L, the greatest original length in that group, and a second pass copies each vector into a fresh list and appends pad_value until its length equals L.

Invariant: after processing a group, the output holds exactly one padded batch per retained group seen so far, in input order, and every vector inside a retained batch has the same length L computed from the original lengths of that group only. Using original lengths (collected before any padding is written) is what makes the padding batch-local: a long vector in one group can never influence L of another group, and copying each vector before appending means the input is never mutated, so a later group cannot observe padding added to an earlier one.

Correctness: the groups partition the input into consecutive runs in order, which matches the grouping rule; the drop test matches the incomplete-batch rule; appending to the end of the copy is right-padding; and the loop emits batches in the order the groups were produced, which is input order.

Edge cases: an empty vectors list produces no iteration and returns []; a batch_size at least as large as the number of vectors produces one group that is incomplete unless the sizes match exactly, so drop_last true can legitimately return []; a group of only empty vectors has L = 0, so the padding loop appends nothing and the empty vectors stay empty; an empty vector inside a group with a longer vector becomes L copies of pad_value; pad_value may be negative, zero, or equal to values already present, and none of those is special-cased; batch_size = 1 makes every group full and self-maximal so no padding occurs.

Time complexity:
O(N + P), where N is the total number of input elements and P is the total number of elements after padding (each vector is scanned for its length and copied once, then padded to its batch maximum).
Space complexity:
O(P) for the returned batches (O(1) auxiliary space beyond the output).