searching in a pandas df that contains ranges

Question

I have a pandas df that contains 2 columns 'start' and 'end' (both are integers). I would like an efficient method to search for rows such that the range that is represented by the row [start,end] contains a specific value.

Two additional notes:

1. It is possible to assume that ranges don't overlap
2. The solution should support a batch mode - that given a list of inputs, the output will be a mapping (dictionary or whatever) to the row indices that contain the matching range.

For example:

       start   end
0      7216    7342
1      7343    7343
2      7344    7471
3      7472    8239
4      8240    8495

and the query of

[7215,7217,7344]

will result in

{7217: 0, 7344: 2}

Thanks!

Answer 1

Brute force solution, could use lots of improvements though.

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

Yields

{7344: 2, 7217: 0}

Answer 2

Selected solution

This new solution could have better performances. But there is a limitation, it will only works if there is no gap between ranges like in the example provided.

# 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}

Previous proposal

# 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}