Python pandasのDataFrameで特定の行から複数行抽出する方法

Question

python初心者です。
csvファイルを読み込んでグラフ化する作業を行おうとしています。
その際に、特定の行から複数行（行数指定して）抽出したいのですがよい方法などありますでしょうか。

下記のようなデータの時に、「抽出対象日付のテーブル」にある日付と一致した
「抽出対象テーブル」の日付から特定の行数を指定して抽出したいです。

抽出対象テーブル

日付	値
2021.05.01	12
2021.05.05	13
2021.05.07	15
2021.05.09	14
2021.05.12	11
2021.05.15	14
2021.05.31	15

抽出対象日付テーブル
| 日付 |
| --- | --
|2021.05.05|
|2021.05.09|

特定行数を4とした場合
下記ふたつが抽出されるイメージです。
①2021.05.05の場合

日付	値
2021.05.05	13
2021.05.07	15
2021.05.09	14
2021.05.12	11

②2021.05.05の場合

日付	値
2021.05.09	14
2021.05.12	11
2021.05.15	14
2021.05.31	15

Answer 1

抽出対象の日付がある行が、何行目かを数値で求めて、スライスで切り出すのできそうです。
何行目かを出すのは、いろいろやり方がありますが、日付列をIndex型にすればget_locが使えるので簡単です。

import pandas as pd

N = 4

df = pd.DataFrame({
    '日付': ['2021.05.01', '2021.05.05', '2021.05.07', '2021.05.09', '2021.05.12', '2021.05.15', '2021.05.31'],
    '値': [12, 13, 15, 14, 11, 14, 15]
})

dt = pd.Index(df['日付'])
for target in ['2021.05.05', '2021.05.09']:
    i = dt.get_loc(target)
    df2 = df.iloc[i:(i + N)]
    print(df2)
#           日付   値
#1  2021.05.05  13
#2  2021.05.07  15
#3  2021.05.09  14
#4  2021.05.12  11
#           日付   値
#3  2021.05.09  14
#4  2021.05.12  11
#5  2021.05.15  14
#6  2021.05.31  15

抽出対象の日付がない場合や、複数ある場合については考慮してないので、そういうケースがある場合は必要に応じて対応してください。

Comments

[WsFkt75O8rrra53]

ご指摘の通り実施したら解決しました。ありがとうございます！

Answer 2

shiftを使ってもできますが、以下の方が簡単でしょう。

>>> max_count = 4
>>> def func(x):
...     global count
...     count -= 1
...     if x == start_date:
...         count = max_count
...     return count > 0
...
>>> print(df)
           日付   値
0  2021.05.01  12
1  2021.05.05  13
2  2021.05.07  15
3  2021.05.09  14
4  2021.05.12  11
5  2021.05.15  14
6  2021.05.31  15
>>>
>>> start_date = '2021.05.05'
>>> count = 0
>>>
>>> print(df[df['日付'].apply(func)])
           日付   値
1  2021.05.05  13
2  2021.05.07  15
3  2021.05.09  14
4  2021.05.12  11
>>>
>>> start_date = '2021.05.09'
>>> count = 0
>>>
>>> print(df[df['日付'].apply(func)])
           日付   値
3  2021.05.09  14
4  2021.05.12  11
5  2021.05.15  14
6  2021.05.31  15

Comments

[WsFkt75O8rrra53]

ご指摘の通り実施したら解決しました。ありがとうございます！
今後の参考にさせていただきます。