diff --git a/quantile_forest/_quantile_forest_fast.pyx b/quantile_forest/_quantile_forest_fast.pyx index 730b3d2..0d2035c 100755 --- a/quantile_forest/_quantile_forest_fast.pyx +++ b/quantile_forest/_quantile_forest_fast.pyx @@ -300,7 +300,7 @@ cpdef vector[double] calc_weighted_quantile( cdef SIZE_t i, j cdef double f cdef double quantile - cdef vector[double] cum_weights + cdef vector[double] cum_weights, sorted_quantile_indices cdef int idx_floor, idx_ceil cdef double p, p_floor, p_ceil cdef double v_floor, v_ceil, frac @@ -317,12 +317,6 @@ cpdef vector[double] calc_weighted_quantile( if not issorted: parallel_qsort_asc(inputs, weights, 0, n_inputs-1) - # Get monotonic sorting of quantiles for efficient calculation. - sorted_quantile_idx = vector[double](n_quantiles) - for i in range((sorted_quantile_idx.size())): - sorted_quantile_idx[i] = i - parallel_qsort_asc(quantiles, sorted_quantile_idx, 0, n_quantiles-1) - cum_weights = vector[double](n_inputs) # Calculate the empirical cumulative distribution function (ECDF). @@ -335,6 +329,12 @@ cpdef vector[double] calc_weighted_quantile( f = cum_weights[n_inputs-1] + 1 - (2*C) + # Get the indices that would sort the quantiles in ascending order. + sorted_quantile_indices = vector[double](n_quantiles) + for i in range((sorted_quantile_indices.size())): + sorted_quantile_indices[i] = i + parallel_qsort_asc(quantiles, sorted_quantile_indices, 0, n_quantiles-1) + out = vector[double](n_quantiles) idx_floor = 0 @@ -344,7 +344,7 @@ cpdef vector[double] calc_weighted_quantile( quantile = quantiles[i] # Assign the output based on the input quantile ordering. - i = sorted_quantile_idx[i] + i = sorted_quantile_indices[i] # Calculate the quantile's proportion of total weight. p = quantile * f + C