Randomly partition array into k segments
Company: Roblox
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: Medium
Interview Round: Technical Screen
##### Question
Given a list of integers and an integer k representing the number of segments, write a function that randomly partitions the list into k non-empty segments and returns them as a list of lists. Explain your approach, discuss its time and space complexity, and describe how you would test the randomness of your solution.
Quick Answer: This question evaluates algorithm design, randomization techniques, and time/space complexity analysis within the Coding & Algorithms domain, focusing on partitioning an array into random non-empty segments.
Given an integer array nums, an integer k (1 <= k <= len(nums)), and a 32-bit unsigned integer seed, return a uniformly random partition of nums into k non-empty contiguous segments. Model the partition as placing k-1 cuts among the n-1 gaps between elements, uniformly over all C(n-1, k-1) choices. For reproducibility, use exactly the following PRNG and sampling method:
- PRNG: 32-bit linear congruential generator (LCG). State is initialized as seed mod 2^32. Each draw sets state = (1664525 * state + 1013904223) mod 2^32 and returns state.
- To draw a uniform integer in [0, m-1], use rejection sampling: let M = 2^32 and limit = (M // m) * m. Repeatedly draw r until r < limit, then return r % m.
- To choose k-1 distinct cut positions uniformly from [1, n-1], use Floyd's algorithm: iterate j from (n-1)-(k-1)+1 to (n-1), pick t uniformly in [1, j]; if t is already chosen, include j; otherwise include t. Sort the selected positions and cut after those indices to form segments.
Constraints
- 1 <= len(nums) <= 100000
- 1 <= k <= len(nums)
- -10^9 <= nums[i] <= 10^9
- 0 <= seed < 2^32
- Output must be k non-empty contiguous subarrays whose concatenation equals nums
Solution
from typing import List
class LCG32:
def __init__(self, seed: int):
self.state = seed & 0xFFFFFFFF
def next32(self) -> int:
# state = (1664525 * state + 1013904223) mod 2^32
self.state = (1664525 * self.state + 1013904223) & 0xFFFFFFFF
return self.state
def randbelow(self, n: int) -> int:
if n <= 0:
raise ValueError("n must be positive")
modulus = 1 << 32
limit = (modulus // n) * n
while True:
r = self.next32()
if r < limit:
return r % n
def randint(self, a: int, b: int) -> int:
# inclusive range [a, b]
return a + self.randbelow(b - a + 1)
def _sample_unique_1_to_N(N: int, m: int, rng: LCG32) -> List[int]:
# Floyd's algorithm: sample m unique integers from [1..N]
selected = set()
if m == 0:
return []
start = N - m + 1
for j in range(start, N + 1):
t = rng.randint(1, j)
if t in selected:
selected.add(j)
else:
selected.add(t)
return sorted(selected)
def random_partition(nums: List[int], k: int, seed: int) -> List[List[int]]:
n = len(nums)
if not (1 <= k <= n):
raise ValueError("k must satisfy 1 <= k <= len(nums)")
rng = LCG32(seed)
if k == 1:
return [nums[:]]
N = n - 1 # number of gaps
m = k - 1 # number of cuts
cuts = _sample_unique_1_to_N(N, m, rng)
res: List[List[int]] = []
prev = 0
for pos in cuts:
res.append(nums[prev:pos])
prev = pos
res.append(nums[prev:])
return res
Explanation
The partition is determined by k-1 cut positions among n-1 available gaps. To ensure uniformity and reproducibility, we use a fixed 32-bit LCG and rejection sampling to obtain unbiased integers in ranges. Floyd's algorithm samples k-1 distinct gap indices uniformly without replacement in O(k) time by iterating j and choosing t uniformly in [1..j], resolving collisions by substituting j. After sorting the cuts, we slice nums at those positions to produce k non-empty contiguous segments. Randomness validation (not required by the function) can be done by running many trials and comparing observed frequencies of cut positions against the expected uniform distribution using statistical tests such as chi-square.
Time complexity: O(n + k). Space complexity: O(k) extra (excluding output).
Hints
- Reduce the problem to choosing k-1 distinct cut positions among n-1 gaps.
- Use a seeded PRNG to make the output deterministic for testing.
- Floyd's algorithm samples without replacement in O(k) time and O(k) space.
- After selecting cuts, sort them and slice the array accordingly.
- To test randomness quality, run many trials and apply frequency checks (e.g., chi-square) over cut positions.