Поиск в панде df, которая содержит диапазоны
У меня есть панды DF, который содержит 2 столбца "начало" и "конец" (оба являются целыми числами). Я хотел бы, чтобы эффективный метод поиска строк, чтобы диапазон, представленный строкой [начало, конец], содержал определенное значение.
Два дополнительных примечания:
- Можно предположить, что диапазоны не перекрываются
- Решение должно поддерживать пакетный режим - при наличии списка входных данных выходные данные будут отображаться (словарь или любой другой) в индексах строк, которые содержат соответствующий диапазон.
Например:
start end
0 7216 7342
1 7343 7343
2 7344 7471
3 7472 8239
4 8240 8495
и запрос
[7215,7217,7344]
приведет к
{7217: 0, 7344: 2}
Спасибо!
2 ответа
Решение грубой силы, может использовать много улучшений, хотя.
df = pd.DataFrame({'start': [7216, 7343, 7344, 7472, 8240],
'end': [7342, 7343, 7471, 8239, 8495]})
search = [7215, 7217, 7344]
res = {}
for i in search:
mask = (df.start <= i) & (df.end >= i)
idx = df[mask].index.values
if len(idx):
res[i] = idx[0]
print res
Урожайность
{7344: 2, 7217: 0}
Выбранное решение
Это новое решение может иметь лучшие характеристики. Но есть ограничение, оно будет работать только в том случае, если между диапазонами нет промежутка, как в приведенном примере.
# Test data
df = pd.DataFrame({'start': [7216, 7343, 7344, 7472, 8240],
'end': [7342, 7343, 7471, 8239, 8495]}, columns=['start','end'])
query = [7215,7217,7344]
# Reshaping the original DataFrame
df = df.reset_index()
df = pd.concat([df['start'], df['end']]).reset_index()
df = df.set_index(0).sort_index()
# Creating a DataFrame with a continuous index
max_range = max(df.index) + 1
min_range = min(df.index)
s = pd.DataFrame(index=range(min_range,max_range))
# Joining them
s = s.join(df)
# Filling the gaps
s = s.fillna(method='backfill')
# Then a simple selection gives the result
s.loc[query,:].dropna().to_dict()['index']
# Result
{7217: 0.0, 7344: 2.0}
Предыдущее предложение
# Test data
df = pd.DataFrame({'start': [7216, 7343, 7344, 7472, 8240],
'end': [7342, 7343, 7471, 8239, 8495]}, columns=['start','end'])
# Constructing a DataFrame containing the query numbers
query = [7215,7217,7344]
result = pd.DataFrame(np.tile(query, (len(df), 1)), columns=query)
# Merging the data and the query
df = pd.concat([df, result], axis=1)
# Making the test
df = df.apply(lambda x: (x >= x['start']) & (x <= x['end']), axis=1).loc[:,query]
# Keeping only values found
df = df[df==True]
df = df.dropna(how='all', axis=(0,1))
# Extracting to the output format
result = df.to_dict('split')
result = dict(zip(result['columns'], result['index']))
# The result
{7217: 0, 7344: 2}