import collections


def map_features_to_example(
    examples, features
):
    """Returns the mapping of feature indices to example index.

    Args:
        examples ([Dataset]): raw datasets.
        features ([Dataset]): tokenized & splited datasets.

    Returns:
        [Dict[List]]: the mapping of feature indices to example index.
    """

    # example과 mapping되는 feature 생성
    example_id_to_index = {k: i for i, k in enumerate(examples["id"])}
    features_per_example = collections.defaultdict(list)
    for i, feature in enumerate(features):
        features_per_example[example_id_to_index[feature["example_id"]]].append(
            i)

    return features_per_example