Как я могу сделать свой код pandas более эффективным?

Question

У меня есть следующий Pandas фрейм данных:

    chr     POS     RS          REF ALT     
1   chr1    981931  rs2465128   A   GA  
2   chr1    982994  rs10267     T   C   
3   chr1    984302  rs9442391   T   C   
4   chr1    987200  rs9803031   C   T   
5   chr1    990280  rs4275402   CT  C   

Я хотел бы иметь еще один столбец со значением «SNP», если длина столбцов «REF» и «ALT» равна равным 1, и значение «INDEL», если какое-либо из них отличается от 1, поэтому результат должен выглядеть следующим образом:

    chr     POS     RS          REF ALT TYPE
1   chr1    981931  rs2465128   A   GA  INDEL
2   chr1    982994  rs10267     T   C   SNP
3   chr1    984302  rs9442391   T   C   SNP
4   chr1    987200  rs9803031   C   T   SNP
5   chr1    990280  rs4275402   CT  C   INDEL

Я написал код, и он работает, но он очень медленный, Мне было интересно, есть ли более эффективный способ сделать это с помощью списков понимания или лямбда-функций.

Мой код

for index, row in table.iterrows():

     if len(row['REF']) == 1 and len(row['ALT']) == 1 :

          table.loc[ index, "TYPE" ] = "SNP"

      else :

          table.loc[ index, "TYPE" ] = "INDEL"

Большое спасибо

Rachael

Answer 1

Используйте Series.str.len для длины и установите новый столбец как numpy.where:

m = (table['REF'].str.len() == 1) & (table['ALT'].str.len() == 1)

table["TYPE"] = np.where(m, "SNP", "INDEL")
print (table)
    chr     POS         RS REF ALT   TYPE
1  chr1  981931  rs2465128   A  GA  INDEL
2  chr1  982994    rs10267   T   C    SNP
3  chr1  984302  rs9442391   T   C    SNP
4  chr1  987200  rs9803031   C   T    SNP
5  chr1  990280  rs4275402  CT   C  INDEL