Python Pandas按条件递增赋值DataFrame列异常问题求助

问题描述 #

我是Python新手，在使用DataFrame时遇到问题。编写了如下Visualize类代码，期望当满足row[x]['lower'] < row[x-1]['upper']条件时，将递增后的max_stage值赋值给stage列（max_stage初始值为0，每次满足条件时自增）。但当前代码运行后，stage列仅返回值1，不符合预期。

原代码：

class Visualize():
    def __init__(self, band, startFreq, endFreq):
        self.band = band
        self.startFreq = startFreq
        self.endFreq = endFreq
        self.max_stage = 0

    # increments max_stage value
    def incrStage(self):
        self.max_stage = self.max_stage+1
        return self.max_stage
    
    
    def getVisualize(self):
        self.band = self.band.sort_values("TX_FREQ")            
        
        self.band['stage'] = np.where(self.band['lower'] < self.band['upper'].shift(), self.incrStage(), 0 )
        
        self.band['difference'] = self.band['lower'] - self.band['upper'].shift()
        
        return self.band

问题原因 #

问题出在np.where的执行逻辑上——它是向量化批量操作，只会在执行时调用一次incrStage()，所以所有满足条件的行都会拿到同一个返回值（第一次调用后的1），而不会每次满足条件就触发一次递增。这就导致stage列所有符合条件的行都是1，而不是你期望的逐行递增的数值。

修复方案 #

方案一：向量化cumsum实现（推荐，效率更高） #

不需要依赖类的max_stage属性，直接通过条件累加生成stage列：

class Visualize():
    def __init__(self, band, startFreq, endFreq):
        self.band = band
        self.startFreq = startFreq
        self.endFreq = endFreq

    def getVisualize(self):
        self.band = self.band.sort_values("TX_FREQ")            
        
        # 生成条件掩码，第一行无前置行，默认不满足条件
        condition = self.band['lower'] < self.band['upper'].shift()
        condition.iloc[0] = False
        
        # 对满足条件的行累加，得到递增的stage值
        self.band['stage'] = condition.cumsum()
        
        self.band['difference'] = self.band['lower'] - self.band['upper'].shift()
        
        return self.band

condition.cumsum()会对满足条件的行逐次累加，每满足一次就加1，正好实现逐行递增的stage效果。

方案二：保留类属性的循环实现（适合理解逻辑，小数据量可用） #

如果一定要用类的max_stage属性，需要逐行遍历DataFrame判断：

class Visualize():
    def __init__(self, band, startFreq, endFreq):
        self.band = band
        self.startFreq = startFreq
        self.endFreq = endFreq
        self.max_stage = 0

    def incrStage(self):
        self.max_stage += 1
        return self.max_stage
    
    def getVisualize(self):
        self.band = self.band.sort_values("TX_FREQ").reset_index(drop=True)
        self.band['stage'] = 0
        
        for idx in range(1, len(self.band)):
            if self.band.loc[idx, 'lower'] < self.band.loc[idx-1, 'upper']:
                self.band.loc[idx, 'stage'] = self.incrStage()
            else:
                # 若不需要重置stage，可保持当前max_stage；否则设0，按需求调整
                self.band.loc[idx, 'stage'] = self.max_stage
        
        self.band['difference'] = self.band['lower'] - self.band['upper'].shift()
        
        return self.band

注意：循环遍历DataFrame在数据量大时效率极低，优先推荐方案一的向量化实现。

内容的提问来源于stack exchange，提问作者Saud Ahmed