def get_chunk_sizes(dataset):
    """
    Given a dataset sorted by data_class
    Return the number of elements of each class, in sorted order
    """
    chunk_sizes = []
    cur_chunk_size = 0
    cur_class = dataset[0]["data_class"]
    allocated_classes = set()
    for seq in dataset:
        assert seq["data_class"] not in allocated_classes, "Data not sorted by class"
        if seq["data_class"] != cur_class:
            allocated_classes.add(cur_class)
            cur_class = seq["data_class"]
            chunk_sizes.append(cur_chunk_size)
            cur_chunk_size = 1
        else:
            cur_chunk_size += 1
    chunk_sizes.append(cur_chunk_size) # Add last chunk since no class change will occur at end
    assert sum(chunk_sizes) == len(dataset)
    return chunk_sizes