python - 基于多列关系合并 Pandas 数据框

Question

假设您有一个区域(开始、结束)坐标的 DataFrame 和另一个可能位于或不属于给定区域内的位置的 DataFrame。例如:

region = pd.DataFrame({'chromosome': [1, 1, 1, 1, 2, 2, 2, 2], 'start': [1000, 2000, 3000, 4000, 1000, 2000, 3000, 4000], 'end': [2000, 3000, 4000, 5000, 2000, 3000, 4000, 5000]})
position = pd.DataFrame({'chromosome': [1, 2, 1, 3, 2, 1, 1], 'BP': [1500, 1100, 10000, 2200, 3300, 400, 5000]})
print region
print position


   chromosome   end  start
0           1  2000   1000
1           1  3000   2000
2           1  4000   3000
3           1  5000   4000
4           2  2000   1000
5           2  3000   2000
6           2  4000   3000
7           2  5000   4000

      BP  chromosome
0   1500           1
1   1100           2
2  10000           1
3   2200           3
4   3300           2
5    400           1
6   5000           1

如果满足以下条件，则位置落在区域内:

position['BP'] >= region['start'] &
position['BP'] <= region['end'] &
position['chromosome'] == region['chromosome']

每个位置保证最多落在一个区域中，尽管它可能不落在任何区域中。

合并这两个数据框的最佳方法是什么，以便它附加额外的列以定位它所在的区域(如果它落在任何区域中)。在这种情况下大致给出以下输出:

      BP  chromosome  start  end
0   1500           1  1000   2000
1   1100           2  1000   2000
2  10000           1  NA     NA
3   2200           3  NA     NA
4   3300           2  3000   4000
5    400           1  NA     NA
6   5000           1  4000   5000

一种方法是编写一个函数来计算我想要的关系，然后使用 DataFrame.apply 方法，如下所示:

def within(pos, regs):
    istrue = (pos.loc['chromosome'] == regs['chromosome']) & (pos.loc['BP'] >= regs['start']) &  (pos.loc['BP'] <= regs['end'])
    if istrue.any():
        ind = regs.index[istrue].values[0]
        return(regs.loc[ind ,['start', 'end']])
    else:
        return(pd.Series([None, None], index=['start', 'end']))

position[['start', 'end']] = position.apply(lambda x: within(x, region), axis=1)
print position

      BP  chromosome  start   end
0   1500           1   1000  2000
1   1100           2   1000  2000
2  10000           1    NaN   NaN
3   2200           3    NaN   NaN
4   3300           2   3000  4000
5    400           1    NaN   NaN
6   5000           1   4000  5000

但我希望有一种比在 O(N) 时间内进行每次比较更优化的方法。谢谢!

Answer 1

一个解决方案是在 chromosome 上进行内连接，排除违规行，然后使用 position 进行左连接:

>>> df = pd.merge(position, region, on='chromosome', how='inner')
>>> idx = (df['BP'] < df['start']) | (df['end'] < df['BP'])  # violating rows
>>> pd.merge(position, df[~idx], on=['BP', 'chromosome'], how='left')
      BP  chromosome   end  start
0   1500           1  2000   1000
1   1100           2  2000   1000
2  10000           1   NaN    NaN
3   2200           3   NaN    NaN
4   3300           2  4000   3000
5    400           1   NaN    NaN
6   5000           1  5000   4000