This is my data frame. If I search for Iowa, the code should return the country name (the USA in this case) and 3 states above it (Hawaii, California, Missouri)-Iowa- and 3 below it (colorado, Alaska, texas in this case). How to do this?

CodePudding user response：

Try this:

import pandas as pd

# prepare the df
c = ['India', '', '', '', '', '', 'USA', '', '', '', '', '', '', '']
s = ['Banglore', 'Pune', 'Delhi', 'Maharasthra', 'Hyderabad', 'Gujarat',
     'Arizona', 'Hawaii', 'California', 'Missouri', 'Iowa', 'Colorado', 'Alaska', 'Texas']

df = pd.DataFrame(c, columns=['Country'])
df['State'] = s

df['Country'][1:6] = 'India'
df['Country'][7:14] = 'USA'

print(df)

def get_states(data, state_name):

    # find the country
    result = data[data['State'] == state_name]
    country_name = list(result['Country'])[0]

    # filter df by country
    all_states = data[data['Country'] == country_name]
    l = list(all_states['State'])
    index = l.index(state_name)
    start = index - 3
    end = index   4
    if index < 3:
       start = 0
    if index   4 > len(l):
        end = len(l)
    final_list = l[start:end]

    return final_list


final_l = get_states(df, 'Iowa')
print(final_l)

CodePudding user response：

Considering your df to be:

In [814]: df
Out[814]: 
   Country        State
0    India     Banglore
1    India         Pune
2    India        Delhi
3    India  Maharasthra
4    India    Hyderabad
5    India      Gujarat
6      USA      Arizona
7      USA       Hawaii
8      USA   California
9      USA     Missouri
10     USA         Iowa
11     USA     Colorado
12     USA       Alaska
13     USA        Texas

Use df.iloc:

In [815]: ix = df[df.State.eq('Iowa')].index[0]

In [817]: res = df.iloc[ix - 3: ix   3]

In [818]: res
Out[818]: 
   Country       State
7      USA      Hawaii
8      USA  California
9      USA    Missouri
10     USA        Iowa
11     USA    Colorado
12     USA      Alaska