如何提高我的熊猫代码效率?

Question

我有以下熊猫数据帧：

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   

如果"REF"和"ALT"列的长度都等于 1，我想要另一列的值为"SNP"，如果它们中的任何一个与 1 不同，则值为"INDEL"，因此输出应如下所示：

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

我已经编写了一些代码，它确实有效，但它非常慢，我想知道是否有更有效的方法可以通过理解列表或 lambda 函数来做到这一点。

我的代码

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"

多谢

瑞 秋

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