重构Python嵌套字典：实现“轴向”层级交换

data_group_test = {
    "Example": {
        "model": {
            "epoch_X": {  # X代表具体的epoch，如epoch1, epoch20
                "dataset_A": {  # A代表具体的dataset，如dataset1, dataset_a
                    # ... 内部数据，例如np.array
                }
            }
        }
    }
}

data_group_test = {
    "Example": {
        "model": {
            "dataset_A": {
                "epoch_X": {
                    # ... 内部数据，例如np.array
                }
            }
        }
    }
}

import json

# 初始数据结构
# 结构：example -> model -> epoch -> dataset -> (data_item)
data_group_tests = {
    "example1": {
        "model1": {
            "epoch1": {
                "epoch1_item": "data_for_epoch1_and_dataset_X",
                "dataset1": {"dataset1_item": "data_for_dataset1_and_epoch1"}
            },
            "epoch2": {
                "epoch2_item": "data_for_epoch2_and_dataset_Y",
                "dataset1": {"dataset1_item": "data_for_dataset1_and_epoch2"}
            }
        }
    }
}

print("--- 原始数据结构 ---")
print(json.dumps(data_group_tests, indent=4))

# 假设我们要针对 "example1" -> "model1" 下的结构进行重构
# 目标：将 "epoch" 和 "dataset" 层级互换

# 步骤 1: 遍历并重构
# 为了通用性，我们可以遍历所有模型，但此处以一个具体路径为例
target_model_path = data_group_tests["example1"]["model1"]

# 假设每个epoch下都有一个或多个dataset，我们需要为每个dataset创建新的结构
# 并且把所有相关epoch的数据都收集到对应的dataset下

# 1. 收集所有 dataset 的名称
all_datasets = set()
for epoch_key, epoch_data in target_model_path.items():
    if isinstance(epoch_data, dict): # 确保是字典，排除如 epoch1_item 这样的直接数据
        for key in epoch_data.keys():
            if key.startswith("dataset"): # 假设 dataset 键以 "dataset" 开头
                all_datasets.add(key)

# 2. 创建新的模型结构，以 dataset 为主键
new_model_structure = {}
for dataset_key in all_datasets:
    new_model_structure[dataset_key] = {}
    for epoch_key, epoch_data in target_model_path.items():
        if isinstance(epoch_data, dict) and dataset_key in epoch_data:
            # 提取 dataset 的内容
            dataset_content = epoch_data[dataset_key]
            # 提取 epoch 自身的内容（除了 dataset 键）
            epoch_specific_content = {k: v for k, v in epoch_data.items() if not k.startswith("dataset")}

            # 将 epoch 的内容放到新的 dataset -> epoch 结构下
            new_model_structure[dataset_key][epoch_key] = {
                **epoch_specific_content, # 合并 epoch 自身的其他数据
                **dataset_content # 合并 dataset 内部的数据
            }

# 3. 更新原始字典
data_group_tests["example1"]["model1"] = new_model_structure

print("\n--- 重构后的数据结构 ---")
print(json.dumps(data_group_tests, indent=4))

# 注意：如果原始结构中 epoch 下除了 dataset 还有其他数据，需要妥善处理。
# 上述代码将 epoch1_item 和 dataset1_item 都合并到了新的 dataset -> epoch 结构下。
# 如果 epoch1_item 是 epoch 独有的，并且 dataset1_item 是 dataset 独有的，
# 那么合并时需要根据实际业务逻辑进行区分。
# 这里为了简化，假设 epoch1_item 也是 epoch 级别的数据，应该跟随 epoch。