import torch


def duplicate_labels(labels, lengths):
    """
    Duplicates labels tensor according to the lengths tensor.
    More specifically, if labels = [l1, l2, l3] and lengths = [n1, n2, n3],
    then the function returns [l1,...(n1 times)...,l1, l2,...(n2 times)...,l2, l3,...(n3 times)...,l3].

    Parameters
    ----------
    labels: torch.Tensor
        1D tensor of labels
    lengths: torch.Tensor
        1D tensor of labels

    Returns
    -------
    labels_duplicated: torch.Tensor
        1D tensor of duplicated labels
    """

    # for each i, duplicate labels[i] lengths[i] times, and concatenate all those.
    labels_duplicated = torch.cat([labels.view(-1)[i].repeat(lengths[i]).view(-1) for i in range(len(lengths))], dim=0)

    return labels_duplicated