Поиск в панде df, которая содержит диапазоны

У меня есть панды DF, который содержит 2 столбца "начало" и "конец" (оба являются целыми числами). Я хотел бы, чтобы эффективный метод поиска строк, чтобы диапазон, представленный строкой [начало, конец], содержал определенное значение.

Два дополнительных примечания:

  1. Можно предположить, что диапазоны не перекрываются
  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}
Другие вопросы по тегам