メモ：sort()にpartial引数を指定すればpartial sortingでちょっと速くなる

R

ふとbase::quantile()と自作関数を比べていると、baseの方が2～3倍くらい速いことに気付いた。

base_quantile <- function(x, probs) {
  # my_quantile()と結果を合わせるためにunname()する
  # typeは、1か3が「〇%の位置の要素を取り出す」感じのやつっぽい
  unname(quantile(x, probs = probs, type = 3))
}

my_quantile <- function(x, probs) {
  sort(x)[floor(length(x) * probs)]
}

x <- runif(1e6)
bench::mark(
  base_quantile(x, probs = 0.3),
  my_quantile(x, probs = 0.3)
)
#> # A tibble: 2 x 10
#>   expression   min  mean  median     max `itr/sec` mem_alloc  n_gc n_itr
#>   <chr>      <bch> <bch> <bch:t> <bch:t>     <dbl> <bch:byt> <dbl> <int>
#> 1 base_quan~  32ms  33ms  32.4ms  34.6ms     30.3     19.3MB     9     3
#> 2 my_quanti~ 100ms 107ms 105.4ms 114.4ms      9.37    11.4MB     2     3
#> # ... with 1 more variable: total_time <bch:tm>

こんなシンプルな実装なのに、なぜ？と思っていつものようにr-sourceを見に行くと、どうもsort()のpartialという引数が胆らしい。

            x <- sort(x, partial =
                      unique(c(1, j[j>0L & j<=n], (j+1)[j>0L & j<n], n))
                      )

(r-source/quantile.R at 5a156a0865362bb8381dcd69ac335f5174a4f60c · wch/r-source · GitHub)

?sort曰く、

If partial is not NULL, it is taken to contain indices of elements of the result which are to be placed in their correct positions in the sorted array by partial sorting.

とのことで、これはpartial sortingというやつらしい。要は、指定した位置の要素が確定したら並べ替えを終わる（という説明であってる？）。 Rのquantile()の場合は「要素を取り出す」わけではなくて線形補間した値を計算したり、Infの存在を考えてるのでややこしくなってるけど、基本的にはこんな感じでよさそう。

my_quantile2 <- function(x, probs) {
  idx <- floor(length(x) * probs)
  sort(x, partial = idx)[idx]
}

x <- runif(1e6)
bench::mark(
  base_quantile(x, probs = 0.3),
  my_quantile(x, probs = 0.3),
  my_quantile2(x, probs = 0.3)
)
#> # A tibble: 3 x 10
#>   expression    min    mean  median     max `itr/sec` mem_alloc  n_gc n_itr
#>   <chr>      <bch:> <bch:t> <bch:t> <bch:t>     <dbl> <bch:byt> <dbl> <int>
#> 1 base_quan~ 34.4ms  35.9ms  35.9ms  37.3ms     27.9     19.1MB     8     3
#> 2 my_quanti~ 99.1ms 112.6ms 107.1ms 131.7ms      8.88    11.4MB     2     3
#> 3 my_quanti~ 23.2ms  26.1ms  24.2ms  30.9ms     38.2     11.4MB     7    10
#> # ... with 1 more variable: total_time <bch:tm>

なるほど速くなった。

quantile()についてはこのスライド、

ソートのアルゴリズムについてはこのQiitaがまとまってそうだった。流し読みしかしてないけど。