Creating dictionary keys with substrings from another column

Question

I have a dataframe with a column that includes the names of boroughs in New York City (Manhattan, Brooklyn, etc.). I want to create another column, 'borough_num', that assigns a number for each borough (Manhattan -> 1, Brooklyn -> 2, Queens -> 3, Staten Island -> 4, Bronx -> 5, Other -> 0).

However, in the Borough column, some rows contain numbers in front of the borough name (e.g. instead of "Bronx" I have "07 Bronx"). Because this "07 Bronx" is still part of the Bronx borough, it should also be assigned to the same value "5" as "Bronx". I therefore need to create a dictionary that assigns the number 5 to a string that contains the word "Bronx". The same for each Borough. Any clue on how to do this? I am new to python!

This is what I had before noticing the cells with the numbers:

df['Borough'] = df['Borough'].fillna('OTHER')
borough_dict = {'MANHATTAN':1, 'BROOKLYN':2, 'QUEENS': 3, 'STATEN ISLAND': 4, 'BRONX': 5, 'OTHER':6}
df['borough_num'] = df['Borough'].apply(lambda x:0 if borough_dict.get(x) == None else borough_dict.get(x))

Answer 1

Since there is a small set of borough names for which to assign the integer codes, it's perfectly acceptable to do it as a series of explicit logical index assignments, such as with some example data below.

Specifically, there is no need in this case to try to encapsulate the borough-to-code mapping into a dict or a helper function or any fancier apply or map operation on the DataFrame.

Just a set of 5 boring, direct logical assignments.

In [13]: df = pandas.DataFrame({
    'Borough': ["Manhattan", "Brooklyn", "Bronx", "07 Bronx", 
                "109 Staten Island", "03 Brooklyn", "04 Queens"], 
    'Value':[1, 2, 3, 4, 5, 6, 7]
})

In [14]: df
Out[14]:
             Borough  Value
0          Manhattan      1
1           Brooklyn      2
2              Bronx      3
3           07 Bronx      4
4  109 Staten Island      5
5        03 Brooklyn      6
6          04 Queens      7

In [15]: df['Borough_num'] = 6  # everything defaults to the 'other' case

In [16]: df.loc[df.Borough.str.contains("Manhattan"), 'Borough_num'] = 1

In [17]: df.loc[df.Borough.str.contains("Brooklyn"), 'Borough_num'] = 2

In [18]: df.loc[df.Borough.str.contains("Queens"), 'Borough_num'] = 3

In [19]: df.loc[df.Borough.str.contains("Staten Island"), 'Borough_num'] = 4

In [20]: df.loc[df.Borough.str.contains("Bronx"), 'Borough_num'] = 5

In [21]: df
Out[21]: 
             Borough  Value  Borough_num
0          Manhattan      1            1
1           Brooklyn      2            2
2              Bronx      3            5
3           07 Bronx      4            5
4  109 Staten Island      5            4
5        03 Brooklyn      6            2
6          04 Queens      7            3

If you want to encapsulate the borough-to-code mapping for whatever reason, you can do it with a simple dict followed by a loop:

In [30]: borough_code = {'Manhattan': 1, 'Brooklyn': 2, 'Queens': 3,
                         'Staten Island': 4, 'Bronx': 5}

In [31]: for borough, code in borough_code.items():
    ...:     df.loc[df.Borough.str.contains(borough), 'Borough_num'] = code

Unless the DataFrame is gigantic, the repeated vectorized calculation of str.contains will be indistinguishable from mapping a function across the column, but will be much easier to understand.

Answer 2

Maybe write a simple helper function:

def find_borough_id(name):
    for k, v in borough_dict.items():
        if k in name:
            return v
    return 0

df['borough_num'] = df['Borough'].apply(find_borough_id)

Answer 3

One solution is to search which key of borough_dict is a substring of x and returns its associated value:

def get_borough_num(x):
  for key, val in borough_dict.items():
    if key in x:
       return val
  return 0
df['borough_num'] = df['Borough'].apply(get_borough_num)

---

An alternative solution is to assume that all rows use as borough_name the borough name or endswith a space + the borough name. With such assumption, you can get the borough_name using:

x.rsplit(' ')[-1]

which returns the string after the last space if the string contains a space or the whole string otherwise:

"Manhattan".rsplit(' ')[-1] => "Manhattan"
"blah Manhattan".rsplit(' ')[-1] => "Manhattan"

so when ends up with:

get_borough_num = lambda x: borough_dict.get(x.rsplit(' ')[-1], 0)
df['borough_num'] = df['Borough'].apply(get_borough_num)

Answer 4

Let's use Pandas str accessor and string functions such as extract, join, upper and python method map.

Given setup from @alexlaval question:

borough_dict = {'MANHATTAN':1, 'BROOKLYN':2, 'QUEENS': 3, 'STATEN ISLAND': 4, 'BRONX': 5, 'OTHER':6}

and setup from @ely

df = pd.DataFrame({
    'Borough': ["Manhattan", "Brooklyn", "Bronx", "07 Bronx", 
                "109 Staten Island", "03 Brooklyn", "04 Queens","Unknown"], 
    'Value':[1, 2, 3, 4, 5, 6, 7, 8]
})

Let's create a regex to extract boroughs from dataframe column:

x = '(' + '|'.join(borough_dict.keys()) + ')'

Now, let's use extract and map to get borough number

df['Borough_number'] = df.Borough.str.upper()\
                         .str.extract(x, expand=False).fillna('OTHER')\
                         .map(borough_dict)

Output:

             Borough  Value  Borough_number
0          Manhattan      1               1
1           Brooklyn      2               2
2              Bronx      3               5
3           07 Bronx      4               5
4  109 Staten Island      5               4
5        03 Brooklyn      6               2
6          04 Queens      7               3
7            Unknown      8               6

Answer 5

A naive way to do it is to define a function to search for a string and if that string is found, then return the intended id.

def borough_id(borough):
    if 'Bronx' in borough:
        return 5
    elif ...
        ...
        ...
    else:
        return None

df['borough_num'] = df['Borough'].apply(lambda x: borough_id(x))

Answer 6

Functional approach

Use a custom function with a generator expression, and make use of the default value of dict.get to return 'OTHER' if no match is found.

Then apply the function via pd.Series.apply.

df = pd.DataFrame({'Borough': ['07 BRONX', '01 MANHATTAN', 'STATEN ISLAND', '12 QUEENS', 'UNKNOWN']})

d = {'MANHATTAN':1, 'BROOKLYN':2, 'QUEENS': 3, 'STATEN ISLAND': 4, 'BRONX': 5, 'OTHER':6}

def map_borough(x, mapping):
    return mapping.get(next((k for k in mapping if x.endswith(k)), None), 'OTHER')

df['borough_num'] = df['Borough'].apply(map_borough, mapping=d)

print(df)

#          Borough borough_num
# 0       07 BRONX           5
# 1   01 MANHATTAN           1
# 2  STATEN ISLAND           4
# 3      12 QUEENS           3
# 4        UNKNOWN       OTHER

Answer 7

Object-oriented approach

You can subclass dict and then use pd.Series.map.

class dict_endswith(dict):
    def __getitem__(self, value):
        key = next((k for k in self.keys() if value.endswith(k)), None)
        return self.get(key)

df = pd.DataFrame({'Borough': ['07 BRONX', '01 MANHATTAN', 'STATEN ISLAND', '12 QUEENS', 'UNKNOWN']})

d = dict_endswith({'MANHATTAN':1, 'BROOKLYN':2, 'QUEENS': 3, 'STATEN ISLAND': 4, 'BRONX': 5, 'OTHER':6})

df['borough_num'] = df['Borough'].map(lambda x: d[x]).fillna('OTHER')

print(df)

#          Borough borough_num
# 0       07 BRONX           5
# 1   01 MANHATTAN           1
# 2  STATEN ISLAND           4
# 3      12 QUEENS           3
# 4        UNKNOWN       OTHER

Answer 8

Loopy solution

If you want to use a loop, then make your code more readable by separating data structure from logic.

In this example, you can iterate your dictionary items sequentially. @ely's version is better for a larger dataframe.

df = pd.DataFrame({'Borough': ['07 BRONX', '01 MANHATTAN', 'STATEN ISLAND', '12 QUEENS', 'UNKNOWN']})

d = {'MANHATTAN':1, 'BROOKLYN':2, 'QUEENS': 3, 'STATEN ISLAND': 4, 'BRONX': 5, 'OTHER':6}

def map_borough(x, mapping):
    for k, v in mapping.items():
        if k in x:
            return v
    else:
        return 'OTHER'

df['borough_num'] = df['Borough'].apply(map_borough, mapping=d)

#          Borough borough_num
# 0       07 BRONX           5
# 1   01 MANHATTAN           1
# 2  STATEN ISLAND           4
# 3      12 QUEENS           3
# 4        UNKNOWN       OTHER