pandasのDataFrameからある特定の複数行をDataFrameのまま取得したい事があったのですが、少しハマったので備忘録として残しておこうと思います。
結論
DataFramedfからリストrowsの値で指定した行を取得する場合はdf.iloc[rows]とすれば良い。
NumPyの場合
まず、NumPyの場合を見てみます。
次の配列を考えます。
import numpy as np
arr = np.arange(100).reshape(-1,10)
print(arr)
結果
[[ 0 1 2 3 4 5 6 7 8 9]
[10 11 12 13 14 15 16 17 18 19]
[20 21 22 23 24 25 26 27 28 29]
[30 31 32 33 34 35 36 37 38 39]
[40 41 42 43 44 45 46 47 48 49]
[50 51 52 53 54 55 56 57 58 59]
[60 61 62 63 64 65 66 67 68 69]
[70 71 72 73 74 75 76 77 78 79]
[80 81 82 83 84 85 86 87 88 89]
[90 91 92 93 94 95 96 97 98 99]]
この配列に対して、指定した複数の行を取得する場合を考えます。
ここでは、次のリストrowsに格納された値と同じ番号の行を取得します。
rows = [0,1,1,2]
NumPyの場合は、普通にインデックスを指定する様に書くだけで取得できます。
print(arr[rows])
結果
[[ 0 1 2 3 4 5 6 7 8 9]
[10 11 12 13 14 15 16 17 18 19]
[10 11 12 13 14 15 16 17 18 19]
[20 21 22 23 24 25 26 27 28 29]]
pandasのDataFrameの場合
同じようにpandasのDataFrameでやってみます。
pandasの場合、NumPyと同じ様にすると行ではなく列が取得されてしまいます。
import pandas as pd
df = pd.DataFrame(arr)
print(df)
print("-"*50)
print(df[rows])
結果
0 1 2 3 4 5 6 7 8 9
0 0 1 2 3 4 5 6 7 8 9
1 10 11 12 13 14 15 16 17 18 19
2 20 21 22 23 24 25 26 27 28 29
3 30 31 32 33 34 35 36 37 38 39
4 40 41 42 43 44 45 46 47 48 49
5 50 51 52 53 54 55 56 57 58 59
6 60 61 62 63 64 65 66 67 68 69
7 70 71 72 73 74 75 76 77 78 79
8 80 81 82 83 84 85 86 87 88 89
9 90 91 92 93 94 95 96 97 98 99
--------------------------------------------------
0 1 1 2
0 0 1 1 2
1 10 11 11 12
2 20 21 21 22
3 30 31 31 32
4 40 41 41 42
5 50 51 51 52
6 60 61 61 62
7 70 71 71 72
8 80 81 81 82
9 90 91 91 92
当然ですが、列名を自分で設定するなどして、rowsの値が列名に含まれない場合はエラーとなります。
df = pd.DataFrame(arr, columns=[f"x{i}" for i in range(10)])
print(df)
print("-"*50)
print(df[rows])
結果
x0 x1 x2 x3 x4 x5 x6 x7 x8 x9
0 0 1 2 3 4 5 6 7 8 9
1 10 11 12 13 14 15 16 17 18 19
2 20 21 22 23 24 25 26 27 28 29
3 30 31 32 33 34 35 36 37 38 39
4 40 41 42 43 44 45 46 47 48 49
5 50 51 52 53 54 55 56 57 58 59
6 60 61 62 63 64 65 66 67 68 69
7 70 71 72 73 74 75 76 77 78 79
8 80 81 82 83 84 85 86 87 88 89
9 90 91 92 93 94 95 96 97 98 99
--------------------------------------------------
---------------------------------------------------------------------------
KeyError Traceback (most recent call last)
（中略）
KeyError: "None of [Int64Index([0, 1, 1, 2], dtype='int64')] are in the [columns]"
単に目的の行を取得するだけならDataFrameをNumPyの配列に変換してしまえばよいですが、DataFrameのまま行を取得したい場合もあると思います。しかし、一度Numpyの配列に変換して、行を取得した後にもう一度DataFrameに変換するのは冗長ではないでしょうか？（また、この方法では行番号が新たに作成されます）
df = pd.DataFrame(arr, columns=[f"x{i}" for i in range(10)])
print(df)
print("-"*50)
print(pd.DataFrame(df.values[rows], columns=df.columns))
結果
x0 x1 x2 x3 x4 x5 x6 x7 x8 x9
0 0 1 2 3 4 5 6 7 8 9
1 10 11 12 13 14 15 16 17 18 19
2 20 21 22 23 24 25 26 27 28 29
3 30 31 32 33 34 35 36 37 38 39
4 40 41 42 43 44 45 46 47 48 49
5 50 51 52 53 54 55 56 57 58 59
6 60 61 62 63 64 65 66 67 68 69
7 70 71 72 73 74 75 76 77 78 79
8 80 81 82 83 84 85 86 87 88 89
9 90 91 92 93 94 95 96 97 98 99
--------------------------------------------------
x0 x1 x2 x3 x4 x5 x6 x7 x8 x9
0 0 1 2 3 4 5 6 7 8 9
1 10 11 12 13 14 15 16 17 18 19
2 10 11 12 13 14 15 16 17 18 19
3 20 21 22 23 24 25 26 27 28 29
そこで色々調べてみたところ、ilocにリストを渡せることに気が付きました。
df = pd.DataFrame(arr, columns=[f"x{i}" for i in range(10)])
print(df)
print("-"*50)
print(df.iloc[rows])
結果
x0 x1 x2 x3 x4 x5 x6 x7 x8 x9
0 0 1 2 3 4 5 6 7 8 9
1 10 11 12 13 14 15 16 17 18 19
2 20 21 22 23 24 25 26 27 28 29
3 30 31 32 33 34 35 36 37 38 39
4 40 41 42 43 44 45 46 47 48 49
5 50 51 52 53 54 55 56 57 58 59
6 60 61 62 63 64 65 66 67 68 69
7 70 71 72 73 74 75 76 77 78 79
8 80 81 82 83 84 85 86 87 88 89
9 90 91 92 93 94 95 96 97 98 99
--------------------------------------------------
x0 x1 x2 x3 x4 x5 x6 x7 x8 x9
0 0 1 2 3 4 5 6 7 8 9
1 10 11 12 13 14 15 16 17 18 19
1 10 11 12 13 14 15 16 17 18 19
2 20 21 22 23 24 25 26 27 28 29
これで、指定の行をDataFrameのまま取得することができました。