<commit_before>#ifndef DKN_HPP_
#define DKN_HPP_
#include "basics.hpp"

#include <vector>
#include <flann/flann.hpp>

template <typename Distance>
std::vector<typename Distance::ResultType> DKN(
        flann::Index<Distance> &index,
        const flann::Matrix<typename Distance::ElementType> &query,
        int k = 3,
        const flann::SearchParams &search_params = flann::SearchParams())
{   /* Get the distances to the k-th nearest neighbor of each element in query
     * using the passed index and search params.
     */
    using std::vector;
    using flann::Matrix;

    typedef typename Distance::ResultType DistanceType;

    // matrices to store the results
    Matrix<int> indices(new int[query.rows*k], query.rows, k);
    Matrix<DistanceType> dists(new DistanceType[query.rows*k], query.rows, k);

    // build the index if necessary
    index.buildIndex();

    // search!
    index.knnSearch(query, indices, dists, k, search_params);

    // get out just the results we want
    vector<DistanceType> dkn;
    dkn.reserve(query.rows);
    for (size_t i = 0; i < query.rows; i++) {
        dkn.push_back(dists[i][k-1]);
    }

    delete[] indices.ptr();
    delete[] dists.ptr();

    return dkn;
}

#endif
<commit_msg>flann::L2 is the *squared* euclidean distance, silly<commit_after>#ifndef DKN_HPP_
#define DKN_HPP_
#include "basics.hpp"

#include <cmath>
#include <vector>
#include <flann/flann.hpp>

template <typename Distance>
std::vector<typename Distance::ResultType> DKN(
        flann::Index<Distance> &index,
        const flann::Matrix<typename Distance::ElementType> &query,
        int k = 3,
        const flann::SearchParams &search_params = flann::SearchParams(),
        bool take_sqrt = true)
{   /* Get the distances to the k-th nearest neighbor of each element in query
     * using the passed index and search params.
     *
     * Because flann::L2 is actually the squared Euclidean distance, this
     * function by default square-roots the results. Pass take_sqrt=false to
     * avoid this.
     * */
    using std::vector;
    using flann::Matrix;

    typedef typename Distance::ResultType DistanceType;

    // matrices to store the results
    Matrix<int> indices(new int[query.rows*k], query.rows, k);
    Matrix<DistanceType> dists(new DistanceType[query.rows*k], query.rows, k);

    // build the index if necessary
    index.buildIndex();

    // search!
    index.knnSearch(query, indices, dists, k, search_params);

    // get out just the results we want
    vector<DistanceType> dkn;
    dkn.reserve(query.rows);
    if (take_sqrt)
        for (size_t i = 0; i < query.rows; i++)
            dkn.push_back(std::sqrt(dists[i][k-1]));
    else
        for (size_t i = 0; i < query.rows; i++)
            dkn.push_back(dists[i][k-1]);

    delete[] indices.ptr();
    delete[] dists.ptr();

    return dkn;
}

#endif

{'fraction_non_alphanumeric': 0.12391228680821441, 'fraction_numerical': 0.003480682213713888, 'mean_word_length': 3.183406113537118, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7157466', 'n_tokens_mistral': 951, 'n_tokens_neox': 898, 'n_words': 260}