CODEWITHSUNDEEP
rsortingdata.table
Amitai
Amitai

Reputation: 891

How to get indices of top k values for each (selected) column in data.table

How to find the indices of the top k (say k=3) values for each column

> dt <- data.table( x = c(1, 1, 3, 1, 3, 1, 1), y = c(1, 2, 1, 2, 2, 1, 1) )
> dt
   x y
1: 1 1
2: 1 2
3: 3 1
4: 1 2
5: 3 2
6: 1 1
7: 1 1

Required output:

> output.1
   x y
1: 1 2
2: 3 4
3: 5 5

Or even better (notice the additional helpful descending sort in x):

> output.2
   var top1 top2 top3
1:   x    3    5    1
2:   y    2    4    5

Having the output would be already a great help.

Upvotes: 4

Views: 505

Answers (3)

jangorecki
jangorecki

Reputation: 16727

If the order of the elements having same rank doesn't matter then this answer would be also valid.
The order information can be extracted from data.table index. 

library(data.table)
dt = data.table(x = c(1, 1, 3, 1, 3, 1, 1), y = c(1, 2, 1, 2, 2, 1, 1))
set2key(dt, x)
set2key(dt, y)

tail.index = function(dt, index, n){
    idx = attr(attr(dt, "index"), index)
    rev(tail(idx, n))
}

tail.index(dt, "__x", 3L)
#[1] 5 3 7
tail.index(dt, "__y", 3L)
#[1] 5 4 2

Upvotes: 3

akrun
akrun

Reputation: 887951

We can use sort (with index.return=TRUE) after looping over the columns of the dataset with lapply

dt[, lapply(.SD, function(x) sort(head(sort(x, 
          decreasing=TRUE, index.return=TRUE)$ix,3)))]
#   x y
#1: 1 2
#2: 3 4
#3: 5 5

Or use order

dt[, lapply(.SD, function(x) sort(head(order(-x),3)))]

Upvotes: 6

JasonAizkalns
JasonAizkalns

Reputation: 20483

Here's a verbose solution which I'm sure undermines the slickness of the data.table package:

dt$idx <- seq.int(1:nrow(dt))

k <- 3

top_x <- dt[order(-x), idx[1:k]]
top_y <- dt[order(-y), idx[1:k]]

dt_top <- data.table(top_x, top_y)

dt_top
#    top_x top_y
# 1:     3     2
# 2:     5     4
# 3:     1     5

Upvotes: 1

Related Questions