q19’s diary

日本語教育能力検定試験、エスペラント検定、HSKなどに言及するblogです。

Matrixの特定の要素を取り出しvectorに並べる

Rの話です。

Problem

Matrixから特定の要素を抜き出してvectorを作ろうとしたところ、なぜかmatrixが生成されて困ってしまいました。
たとえば、次のようなmatrixを考えます。

> mat <- matrix(1:100, 10, 10)
> mat
      [,1] [,2] [,3] [,4] [,5] [,6] [,7] [,8] [,9] [,10]
 [1,]    1   11   21   31   41   51   61   71   81    91
 [2,]    2   12   22   32   42   52   62   72   82    92
 [3,]    3   13   23   33   43   53   63   73   83    93
 [4,]    4   14   24   34   44   54   64   74   84    94
 [5,]    5   15   25   35   45   55   65   75   85    95
 [6,]    6   16   26   36   46   56   66   76   86    96
 [7,]    7   17   27   37   47   57   67   77   87    97
 [8,]    8   18   28   38   48   58   68   78   88    98
 [9,]    9   19   29   39   49   59   69   79   89    99
[10,]   10   20   30   40   50   60   70   80   90   100

このmatの(1, 4), (2, 5), (3, 6)成分を取り出してvectorを作ろうと考えたのですが、vectorではなくmatrixが得られてしまいました。

> r <- c(1, 2, 3)
> c <- c(4, 5, 6)
> mat[r, c]
     [,1] [,2] [,3]
[1,]   31   41   51
[2,]   32   42   52
[3,]   33   43   53

Solution

調べると、cbind()という人がいい感じにやってくれることがわかりました。

> mat[cbind(r, c)]
[1] 31 42 53

すぐ思いつくような方法として、for文で地道に生成する方法

> dim <- length(r)
> vec <- numeric(dim)
> for(k in 1:dim)
+   vec[k] <- mat[r[k], c[k]]
> vec
[1] 31 42 53

や、diag()によりmatrixの対角成分だけ取り出す方法

> diag(mat[r, c])
[1] 31 42 53

がありますが、数万〜数千万次元の巨大なvectorを作りたかったので、計算時間の観点からfor文は避けたいという気持ちがありました。
実際、計算時間を計測してみると、cbind()の方法は0.035秒、for文の方法は1.6秒かかり、diag()の方法に至っては、実行するとRがフリーズしてしまい、最後まで計算を行えませんでした。

> r <- round(runif(10^6, min = 1, max = 10)) # 適当に10^6コ自然数生成
> c <- round(runif(10^6, min = 1, max = 10)) # 上に同じ
> dim <- length(r) # 10^6
> vec <- numeric(dim)

> start <- Sys.time()
> vec <- mat[cbind(r, c)]
> end <- Sys.time()
> end - start
Time difference of 0.03506589 secs

> start <- Sys.time()
> for(k in 1:dim)
+   vec[k] <- mat[r[k], c[k]]
> end <- Sys.time()
> end - start
Time difference of 1.6093 secs

> start <- Sys.time()
> vec <- diag(mat[r, c]) # ここでRがフリーズする

フリーズの原因は巨大な行列を生成したことっぽいですが、あんまり興味がわかないのでこの辺で。