Label Distribution by Segment

Task

Implement label_distribution_by_segment(segment_ids: np.ndarray, labels: np.ndarray, num_segments: int, num_labels: int) -> np.ndarray. Debug and repair the provided near-complete implementation so it returns the exact empirical label distribution for every segment.

Requirements

• The starter implementation runs, but contains several independent logic bugs. Repair it without changing the function signature.
• segment_ids and labels are one-dimensional integer arrays of equal length. num_segments and num_labels are positive integers.
• Every segment ID is in [0, num_segments), and every label ID is in [0, num_labels). IDs need not all appear.
• Return a float64 array with shape (num_segments, num_labels). Row s represents segment ID s, and column k represents label ID k; both axes are in increasing numeric ID order.
• For a non-empty segment s, output[s, k] is count(segment_ids == s and labels == k) divided by count(segment_ids == s). The denominator is local to that segment.
• A segment with no observations has an all-zero row. A label absent from a non-empty segment has probability 0; do not add smoothing or renormalize missing combinations.
• Every observation, including the first and last, must be counted exactly once. Counts are exact-label counts, not cumulative counts.
• Empty input is valid and returns an all-zero array of the requested shape. Do not mutate either input.

Example

label_distribution_by_segment(
    np.array([0, 0, 0, 1, 1, 1]),
    np.array([0, 1, 1, 0, 0, 2]),
    2,
    3,
)
# array([[0.33333333, 0.66666667, 0.0],
#        [0.66666667, 0.0, 0.33333333]])