import torch


def precision_k(labels, scores, k=5, average=False, largest=False):
    """
    Computes the precision@k for the specified values of k
    or average precision@k if attribute average is True
    :param labels: pytorch.Tensor of shape (n,) containing labels of each value
    :param scores: pytorch.Tensor of shape (n,) containing values
    :param k: number of items that will be taken
    :param average: if True calculate ap@k metric else calculate precision@k metric
    :param largest: if True take k items with the largest value
                    else take k items with the smallest value
    :return: precision@k or ap@k metric
    """

    topk, indices = scores.topk(k, largest=largest)
    topk_target = labels[indices]

    if average:
        cummulative_sum = topk_target.cumsum(dim=0).float()
        ks = torch.arange(1, k + 1, dtype=torch.float32)
        return (cummulative_sum / ks).mean().item()
    else:
        return topk_target.float().mean().item()