def _count_probabilities(primary_weights, secondary_weights, knn):
    """Computes probabilities of all points to be chosen as the second point in pair
    :type primary_weights: numpy.array, shape = [n_samples],
        the first event is generated according to these weights
    :type secondary_weights: numpy.array, shape = [n_samples],
        the second event is chosen between knn of first according to this weights
    :type knn: dict, {event_id: list of neighbors ids}.
    :rtype: numpy.array, shape = [n_samples], the probabilities
    """
    primary_probabilities = primary_weights / numpy.sum(primary_weights)
    secondary_weights = numpy.array(secondary_weights)
    knn_probabilities = numpy.take(secondary_weights, knn)
    knn_probabilities /= knn_probabilities.sum(axis=1, keepdims=True)

    secondary_probabilities = numpy.bincount(knn.flatten(), weights=knn_probabilities.flatten())
    return primary_probabilities * secondary_probabilities
dataset/ETHPy150Open anaderi/lhcb_trigger_ml/hep_ml/toymc.py/_count_probabilities/original