How can I add a new column based on two dataframes and conditions for big data?

user_18337640 0
Viewed 33

How can I add a new column based on two dataframes and conditions?

For example, if df2['x'] is between df1['x']±10 and df2['y'] is between df1['y']±10 and for same year, give 1 otherwise 0.

The codes below worked with relatively small data. But when I used it for big data (df1=210.548 rows, df2= 301,661 rows), I got error message: "The kernel appears to have died. It will restart automatically".

Are there any other ways to solve it? You can find sample dataframes and expected output in this question: How can I add a new column based on two dataframes and conditions

merged = df1.merge(df2, on="year", how="left", suffixes=("1", "2"))
merged["f"] = (
    (merged.x2 >= merged.x1 - 10)
    & (merged.x2 < merged.x1 + 10)
    & (merged.y2 >= merged.y1 - 10)
    & (merged.y2 < merged.y1 + 10)
)
dataset = (
    merged.groupby(["x1", "y1", "year1", "a", "b", "f"])
    .f.any()
    .astype(int)
    .reset_index()
)
jupyter-notebook pandas
1 Answers

Given the following dataframes:

import pandas as pd

df1 = (
    pd.DataFrame(
        {
            "x": [40.1, 50.1, 60.1, 70.1, 80.1, 90.1, 0, 300.1],
            "y": [100.1, 110.1, 120.1, 130.1, 140.1, 150.1, 160.1, 400.1],
            "year": [2000, 2000, 2001, 2001, 2003, 2003, 2003, 2004],
        }
    )
    .sort_values(["year", "x", "y"])
    .reset_index(drop=True)
)
print(df1)
# Output
       x      y  year
0   40.1  100.1  2000
1   50.1  110.1  2000
2   60.1  120.1  2001
3   70.1  130.1  2001
4    0.0  160.1  2003
5   80.1  140.1  2003
6   90.1  150.1  2003
7  300.1  400.1  2004

df2 = (
    pd.DataFrame(
        {
            "x": [92.2, 30.1, 82.6, 51.1, 39.4, 10.1, 0, 299.1],
            "y": [149.3, 100.1, 139.4, 111.1, 100.8, 180.1, 0, 402.5],
            "year": [1950, 1951, 1952, 2000, 2000, 1954, 1955, 2004],
        }
    )
    .sort_values(["year", "x", "y"])
    .reset_index(drop=True)
)
print(df2)
       x      y  year
0   92.2  149.3  1950
1   30.1  100.1  1951
2   82.6  139.4  1952
3   10.1  180.1  1954
4    0.0    0.0  1955
5   39.4  100.8  2000
6   51.1  111.1  2000
7  299.1  402.5  2004

Here is what your code, given minor fixes, outputs:

      x1     y1  year  f
0   40.1  100.1  2000  1
1   50.1  110.1  2000  1
2   60.1  120.1  2001  0
3   70.1  130.1  2001  0
4    0.0  160.1  2003  0
5   80.1  140.1  2003  0
6   90.1  150.1  2003  0
7  300.1  400.1  2004  1

But you can achieve the same result without expensive merging and groupby operations, like this:

# Find common years between df1 and df2
year_match = set(df1["year"]).intersection(df2["year"])

# Slice both dataframes on common years and align indexes
s1 = df1[df1["year"].isin(year_match)]
s2 = df2[df2["year"].isin(year_match)].set_index(s1.index)

# Initialize new "f" column
df1["f"] = 0

# Update "f" column according to given conditions
mask = (
    (s2.x >= s1.x - 10) & (s2.x < s1.x + 10) & (s2.y >= s1.y - 10) & (s2.y < s1.y + 10)
)
df1.loc[mask.index, "f"] = 1

So that:

print(df1)
# Output
       x      y  year  f
0   40.1  100.1  2000  1
1   50.1  110.1  2000  1
2   60.1  120.1  2001  0
3   70.1  130.1  2001  0
4    0.0  160.1  2003  0
5   80.1  140.1  2003  0
6   90.1  150.1  2003  0
7  300.1  400.1  2004  1
user_112460560
Related