pandas - нарезка str (строки) на основе int в другом столбце

Question

У меня есть df:

   colA    colB
0 'abcde'    4
1 'abcde'    2
2 'abcde'    1
3  np.nan   np.nan
4 'wxyz'     3
5 'wxyz'     2

Я хотел бы иметь возможность удалять первые X символов из colA на основе значения в colB и возвращать значение в новый столбец C как показано ниже.

   colA    colB     colC
0 'abcde'    4      'e'
1 'abcde'    2      'cde'
2 'abcde'    1      'bcde'
3  np.nan   np.nan  np.nan
4 'wxyz'     3      'z'
5 'wxyz'     2      'yz'

Я пробовал некоторые .apply лямбда здесь 1 с .str [x:], но столкнулся с проблемой сохранения его обратно из-за нулевых значений в других строках.

Любая помощь высоко ценится!

Answer 1

Вы можете создать пользовательскую функцию для возврата пропущенных значений, если индексация не удалась:

def f(a, b):
    try:
        return a[int(b):]
    except:
        return np.nan

df['colC'] = [f(a,b) for a, b in zip(df['colA'], df['colB'])]

Или:

df['colC'] = df.apply(lambda x: f(x['colA'], x['colB']), axis=1)

---

print (df)
    colA  colB  colC
0  abcde   4.0     e
1  abcde   2.0   cde
2  abcde   1.0  bcde
3    NaN   NaN   NaN
4   wxyz   3.0     z
5   wxyz   2.0    yz

Еще одна идея с тестом на отсутствие пропущенных значений:

df['colC'] = [a[int(b):] if pd.notna(a) and pd.notna(b) 
                         else np.nan 
                         for a, b in zip(df['colA'], df['colB'])]
print (df)
    colA  colB  colC
0  abcde   4.0     e
1  abcde   2.0   cde
2  abcde   1.0  bcde
3    NaN   NaN   NaN
4   wxyz   3.0     z
5   wxyz   2.0    yz

Answer 2

ответ Джезраэля , вероятно, лучше и удобнее для чтения, но если вы хотите сделать это в oneliner, вы можете использовать df.dropna(). Несмотря на страшное имя, он не изменит ваш оригинальный DataFrame, если вы не вызовете его с параметром inplace=True.

>>> df['colC'] = df.dropna().apply(lambda x: x[0][int(x[1]):], axis=1)

>>> print(df)
    colA    colB    colC
0   abcde   4.0     e
1   abcde   2.0     cde
2   abcde   1.0     bcde
3   NaN     NaN     NaN
4   wxyz    3.0     z
5   wxyz    2.0     yz