Numpy: создание последовательности переменной длины из фрейма данных pandas

Question

Допустим, у меня есть следующий фрейм данных:

df_raw = pd.DataFrame({"person_id": [101, 101, 102, 102, 102, 103], "date": [0, 5, 0, 7, 11, 0], "val1": [99, 11, 22, 33, 44, 22], "val2": [77, 88, 22, 66, 55, 33]})

Чего я хочу добиться, так это создать 3-мерный массив NumPy так, чтобы в результате получилось следующее:

np_pros = np.array([[[0, 99, 77], [5, 11, 88]], [[0, 22, 22], [7, 33, 66], [11, 44, 55]], [[0, 22, 33]]])

Другими словами, трехмерный массив должен иметь следующую форму [unique_ids, None, feature_size]. В моем случае число unique_ids равно 3, feature size равно 3 (все столбцы, кроме person_id), а столбец y имеет переменную длину и указывает количество измерений для person_id.

Мне хорошо известно, что я могу создать массив np.zeros((unique_ids, max_num_features, feature_size)), заполнить его, а затем удалить ненужные элементы, но я хочу что-то быстрее. Причина в том, что мой фактический фрейм данных огромен (примерно [50000, 455]), что приведет к массиву с массой примерно [12500, 200, 455].

Ждем ваших ответов!

Answer 1

Вы можете использовать groupby :

import pandas as pd

df_raw = pd.DataFrame({"person_id": [101, 101, 102, 102, 102, 103], "date": [0, 5, 0, 7, 11, 0], "val1": [99, 11, 22, 33, 44, 22], "val2": [77, 88, 22, 66, 55, 33]})

result = [group.values for _, group in df_raw.groupby('person_id')[['date', 'val1', 'val2']]]
print(result)

выход

[array([[  0, 101,  99,  77],
       [  5, 101,  11,  88]]), array([[  0, 102,  22,  22],
       [  7, 102,  33,  66],
       [ 11, 102,  44,  55]]), array([[  0, 103,  22,  33]])]

Answer 2

Другое решение с xarray

---

Давайте создадим измерение, подразумеваемое двойственностью person_id

>>> df['newdim'] = df.person_id.duplicated()
>>> df.newdim    = df.groupby('person_id').newdim.cumsum()
>>> df           = df.set_index(["newdim", "person_id"])
>>> df
                  date  val1  val2
newdim person_id                  
0.0    101           0    99    77
1.0    101           5    11    88
0.0    102           0    22    22
1.0    102           7    33    66
2.0    102          11    44    55
0.0    103           0    22    33

Для удобства чтения мы можем захотеть превратить df в xarray.Dataset -объект

>>> xa = df.to_xarray()
>>> xa
<xarray.Dataset>
Dimensions:    (newdim: 3, person_id: 3)
Coordinates:
  * newdim     (newdim) float64 0.0 1.0 2.0
  * person_id  (person_id) int64 101 102 103
Data variables:
    date       (newdim, person_id) float64 0.0 0.0 0.0 5.0 7.0 nan nan 11.0 nan
    val1       (newdim, person_id) float64 99.0 22.0 22.0 11.0 33.0 nan nan ...
    val2       (newdim, person_id) float64 77.0 22.0 33.0 88.0 66.0 nan nan ...

, а затем в размерно-здоровых массив Numpy

>>> ar = xa.to_array().T.values
>>> ar
array([[[ 0., 99., 77.],
        [ 5., 11., 88.],
        [nan, nan, nan]],

       [[ 0., 22., 22.],
        [ 7., 33., 66.],
        [11., 44., 55.]],

       [[ 0., 22., 33.],
        [nan, nan, nan],
        [nan, nan, nan]]])

Обратите внимание, что nan -значения были введены с помощью принуждения.

Answer 3

Вот один из способов сделать это:

ix = np.flatnonzero(df1.person_id != df1.person_id.shift(1))
np.split(df1.drop('person_id', axis=1).values, ix[1:])

[array([[ 0, 99, 77],
        [ 5, 11, 88]], dtype=int64), 
 array([[ 0, 22, 22],
        [ 7, 33, 66],
        [11, 44, 55]], dtype=int64), 
 array([[ 0, 22, 33]], dtype=int64)]

Подробности

Используйте np.flatnonzero после сравнения df1 со смещенной версией самого себя (pd.shift), чтобы получить индексы, где происходят изменения в person_id:

ix = np.flatnonzero(df1.person_id != df1.person_id.shift(1))
#array([0, 2, 5])

Используйте np.split, чтобы разбить интересующие столбцы информационного кадра в соответствии с полученным индексом:

np.split(df1.drop('person_id', axis=1).values, ix[1:])

[array([[ 0, 99, 77],
        [ 5, 11, 88]], dtype=int64), 
 array([[ 0, 22, 22],
        [ 7, 33, 66],
        [11, 44, 55]], dtype=int64), 
 array([[ 0, 22, 33]], dtype=int64)]