# 续写同一文件：training_project/src/readingtorch/data.py
def split_indices(y, task, seed):
    # 先划分原始样本，图形按类别分别打乱，维持三组中的类别比例。
    generator = torch.Generator().manual_seed(seed)
    groups = [torch.arange(len(y))]
    if task == "images":
        groups = [torch.where(y == label)[0] for label in range(3)]
    parts = [[], [], []]
    for group in groups:
        group = group[torch.randperm(len(group), generator=generator)]
        first, second = int(len(group) * 0.6), int(len(group) * 0.8)
        for bucket, chunk in zip(
            parts, [group[:first], group[first:second], group[second:]]
        ):
            bucket.append(chunk)
    return {
        name: torch.cat(items) for name, items in zip(["train", "val", "test"], parts)
    }


@dataclass
class Bundle:
    x: torch.Tensor
    y: torch.Tensor
    splits: dict
    preprocess: dict
    source_hash: str