From df05b9d06090312b432489790d8e7d12c47b7e22 Mon Sep 17 00:00:00 2001 From: Benjamin Schwendinger Date: Mon, 15 Dec 2025 12:26:08 +0100 Subject: [PATCH 1/2] add floyd-rivest and parallelization to gmedian --- src/data.table.h | 3 ++ src/gsumm.c | 69 ++++++++++++++++++++++++++++++--------------- src/quickselect.c | 71 +++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 121 insertions(+), 22 deletions(-) diff --git a/src/data.table.h b/src/data.table.h index ebfa1232c6..48f1d346dd 100644 --- a/src/data.table.h +++ b/src/data.table.h @@ -229,6 +229,9 @@ SEXP bmerge(SEXP idt, SEXP xdt, SEXP icolsArg, SEXP xcolsArg, double dquickselect(double *x, int n); double iquickselect(int *x, int n); double i64quickselect(int64_t *x, int n); +double dfloyd_rivest(double *x, int n); +double ifloyd_rivest(int *x, int n); +double i64floyd_rivest(int64_t *x, int n); // fread.c double wallclock(void); diff --git a/src/gsumm.c b/src/gsumm.c index 5970f59194..3cf8d68a34 100644 --- a/src/gsumm.c +++ b/src/gsumm.c @@ -880,35 +880,60 @@ SEXP gmedian(SEXP x, SEXP narmArg) { const bool nosubset = irowslen==-1; switch(TYPEOF(x)) { case REALSXP: { - double *subd = REAL(PROTECT(allocVector(REALSXP, maxgrpn))); // allocate once upfront and reuse int64_t *xi64 = (int64_t *)REAL(x); double *xd = REAL(x); - for (int i=0; i 100 ? i64floyd_rivest((int64_t *)thread_subd, thisgrpsize) : i64quickselect((int64_t *)thread_subd, thisgrpsize)) : + (thisgrpsize > 100 ? dfloyd_rivest(thread_subd, thisgrpsize) : dquickselect(thread_subd, thisgrpsize)); + } } - thisgrpsize -= nacount; // all-NA is returned as NA_REAL via n==0 case inside *quickselect - ansd[i] = (nacount && !narm) ? NA_REAL : (isInt64 ? i64quickselect((void *)subd, thisgrpsize) : dquickselect(subd, thisgrpsize)); + free(thread_subd); }} break; case LGLSXP: case INTSXP: { - int *subi = INTEGER(PROTECT(allocVector(INTSXP, maxgrpn))); int *xi = INTEGER(x); - for (int i=0; i 100 ? ifloyd_rivest(thread_subi, validsize) : iquickselect(thread_subi, validsize); + } } - ansd[i] = (nacount && !narm) ? NA_REAL : iquickselect(subi, thisgrpsize-nacount); + free(thread_subi); }} break; default: @@ -916,7 +941,7 @@ SEXP gmedian(SEXP x, SEXP narmArg) { } if (!isInt64) copyMostAttrib(x, ans); // else the integer64 class needs to be dropped since double is always returned by gmedian - UNPROTECT(2); // ans, subd|subi + UNPROTECT(1); // ans return ans; } diff --git a/src/quickselect.c b/src/quickselect.c index 5277c0d65a..4ca909b140 100644 --- a/src/quickselect.c +++ b/src/quickselect.c @@ -71,3 +71,74 @@ double i64quickselect(int64_t *x, int n) int64_t a, b; BODY(i64swap); } + +// Floyd-Rivest selection algorithm + +#undef FLOYD_RIVEST_BODY +#define FLOYD_RIVEST_BODY(SWAP, TYPE) \ + if (n == 0) return NA_REAL; \ + unsigned long med = n / 2 - (n % 2 == 0); \ + unsigned long l = 0, ir = n - 1; \ + while (ir > l) { \ + if (ir - l > 600) { \ + unsigned long nn = ir - l + 1; \ + unsigned long i = med - l + 1; \ + double z = log((double)nn); \ + double s = 0.5 * exp(2.0*z/3.0); \ + double sd = 0.5 * sqrt(z*s*(nn-s)/nn); \ + if (i < nn/2) sd = -sd; \ + unsigned long newL = (unsigned long)(MAX((long)l, (long)(med - i*s/nn + sd))); \ + unsigned long newR = (unsigned long)(MIN((long)ir, (long)(med + (nn-i)*s/nn + sd))); \ + l = newL; \ + ir = newR; \ + } \ + TYPE pivot = x[med]; \ + unsigned long i = l, j = ir; \ + SWAP(x + l, x + med); \ + if (x[ir] > pivot) { \ + SWAP(x + ir, x + l); \ + pivot = x[l]; \ + } \ + while (i < j) { \ + SWAP(x + i, x + j); \ + i++; j--; \ + while (x[i] < pivot) i++; \ + while (x[j] > pivot) j--; \ + } \ + if (x[l] == pivot) { \ + SWAP(x + l, x + j); \ + } else { \ + j++; \ + SWAP(x + j, x + ir); \ + } \ + if (j <= med) l = j + 1; \ + if (med <= j) ir = j - 1; \ + } \ + a = x[med]; \ + if (n % 2 == 1) { \ + return (double)a; \ + } else { \ + b = x[med + 1]; \ + for (unsigned long i = med + 2; i < n; i++) { \ + if (x[i] < b) b = x[i]; \ + } \ + return ((double)a + (double)b) / 2.0; \ + } + +double dfloyd_rivest(double *x, int n) +{ + double a, b; + FLOYD_RIVEST_BODY(dswap, double); +} + +double ifloyd_rivest(int *x, int n) +{ + int a, b; + FLOYD_RIVEST_BODY(iswap, int); +} + +double i64floyd_rivest(int64_t *x, int n) +{ + int64_t a, b; + FLOYD_RIVEST_BODY(i64swap, int64_t); +} From 36742ea61274eeb31e967279a2229cee2762b578 Mon Sep 17 00:00:00 2001 From: Benjamin Schwendinger <52290390+ben-schwen@users.noreply.github.com> Date: Tue, 16 Dec 2025 21:14:13 +0100 Subject: [PATCH 2/2] Update src/quickselect.c Co-authored-by: Michael Chirico --- src/quickselect.c | 1 + 1 file changed, 1 insertion(+) diff --git a/src/quickselect.c b/src/quickselect.c index 4ca909b140..67dfe34e3e 100644 --- a/src/quickselect.c +++ b/src/quickselect.c @@ -73,6 +73,7 @@ double i64quickselect(int64_t *x, int n) } // Floyd-Rivest selection algorithm +// https://en.wikipedia.org/wiki/Floyd%E2%80%93Rivest_algorithm #undef FLOYD_RIVEST_BODY #define FLOYD_RIVEST_BODY(SWAP, TYPE) \