pandasのrolling()に条件を付け加えて要素を抽出し、元のDataFrameに結合したい

Question

前提

下に用意したデータで、5つずつ区間をシフトさせ、それぞれについて男性(M)であれば、その時の人数(num)の合計を求めて、新しい列(num_M)として元のDataFrameに結合したいです。女性(F)についても同じです。for文で求めようとしたのですが、欲しい結果が得られせんし、読みづらいので簡潔なコードにしたいです。(rolling()を使わなければいけないというわけではないです)。

time	sex	num
2018-07-01 10:00:00	M	1
2018-07-01 10:00:00	M	1
2018-07-01 10:00:00	M	5
2018-07-01 10:00:00	F	3
2018-07-01 10:00:02	F	5
2018-07-01 10:00:03	M	1
2018-07-01 10:00:03	F	3
2018-07-01 10:00:03	F	4
2018-07-01 10:00:05	F	2
2018-07-01 10:00:10	M	1

欲しい結果

time	sex	num	num_M	num_F
2018-07-01 10:00:00	M	1	NaN	NaN
2018-07-01 10:00:00	M	1	NaN	NaN
2018-07-01 10:00:00	M	5	NaN	NaN
2018-07-01 10:00:00	F	3	NaN	NaN
2018-07-01 10:00:02	F	5	7	8
2018-07-01 10:00:03	M	1	7	8
2018-07-01 10:00:03	F	3	6	11
2018-07-01 10:00:03	F	4	1	15
2018-07-01 10:00:05	F	2	1	14
2018-07-01 10:00:10	M	1	2	9

該当のソースコード

import numpy as np
import pandas as pd

# 元データ
index = ['2018-07-01 10:00:00', '2018-07-01 10:00:00', '2018-07-01 10:00:00', '2018-07-01 10:00:00', '2018-07-01 10:00:02', '2018-07-01 10:00:03', '2018-07-01 10:00:03', '2018-07-01 10:00:03', '2018-07-01 10:00:05', '2018-07-01 10:00:10']
columns = ['sex']

df = pd.DataFrame(list('MMMFFMFFFM'), index=index, columns=columns)

# 列(num)を追加
df['num'] = [1, 1, 5, 3, 5, 1, 3, 4, 2, 1]
df

# 列(num_M)を追加
df['num_M'] = np.nan

for i in range(len(df)-5):
    df.iloc[i+4, df.columns.get_loc('num_M')] = df.iloc[i:i+5, df.columns.get_loc('num')][df.iloc[i:i+5, :]['sex']=='M'].sum()
df

最終的に数万行のDataFrameに対して処理を実行したいので、for文を使うと処理速度にかなり影響が出るのではと思っています。何かいいアイデアやヒントがあれば教えていただけると嬉しいです。

Answer 1

df['num_M'] = df['num'].where(df.sex=='M',0).rolling(5).sum()
df['num_F'] = df['num'].where(df.sex=='F',0).rolling(5).sum()

でどうでしょうか。

Comments

[退会済みユーザー]

magichanさん、回答有り難うございます。

回答にあるコードでほしい結果を得ることができました。whereを使って、元のindexをそのままに条件にあわない箇所を0にするのは目からウロコでした。ベストアンサーにさせていただきます。ありがとうございましたm(_ _)m