modifying nans position in the dataframe

user_19643955 0
Viewed 30

I'm hoping I can explain this well. I have this df with 2 clumns: group and numbers. I'm trying to get that np.nan and pop it into it's new group.

def check_for_nan():
    # for example let's say my new value is 14.5
    new_nan_value=14.5
    data = {"group:" : [-1,0,1,2,3],
            'numbers': [[np.nan], [11, 12], [14, 15], [16, 17], [18, 19]],
            }
    df = pd.DataFrame(data=data)


    # *** add some code ***


    # I created a new dataframe to visually show how it should look like but we would want to operate only on the same df from above 
    data_2 = {"group" : [0,1,2,3],
            'numbers': [[11, 12], [14,np.nan, 15], [16, 17], [18, 19]],
            }
    df_2 = pd.DataFrame(data=data_2)
    # should return the new group number where the nan would live
    return data_2["group"][1]

Output:

current:

   group:   numbers
0      -1     [nan]
1       0  [11, 12]
2       1  [14, 15]
3       2  [16, 17]
4       3  [18, 19]

Desired output when new_nan_value =14.5

   group        numbers
0      0       [11, 12]
1      1  [14, nan, 15]
2      2       [16, 17]
3      3       [18, 19]

return 1
pandas dataframe nan arrays python
1 Answers

With the dataframe you provided:

import pandas as pd

df = pd.DataFrame(
    {
        "group": [-1, 0, 1, 2, 3],
        "numbers": [[pd.NA], [11, 12], [14, 15], [16, 17], [18, 19]],
    }
)

new_nan_value = 14.5

Here is one way to do it:

def move_nan(df, new_nan_value):
    """Helper function.

    Args:
        df: input dataframe.
        new_nan_value: insertion value.

    Returns:
        Dataframe with nan value at insertion point, new group.

    """

    # Reshape dataframe along row axis
    df = df.explode("numbers").dropna().reset_index(drop=True)

    # Insert new row
    insert_pos = df.loc[df["numbers"] < new_nan_value, "numbers"].index[-1] + 1
    df = pd.concat(
        [
            df.loc[: insert_pos - 1, :],
            pd.DataFrame({"group": [pd.NA], "numbers": pd.NA}, index=[insert_pos]),
            df.loc[insert_pos:, :],
        ]
    )
    df["group"] = df["group"].fillna(method="bfill")

    # Find new group value
    new_group = df.loc[df["numbers"].isna(), "group"].values[0]

    # Groupby and reshape dataframe along column axis
    df = df.groupby("group").agg(list).reset_index(drop=False)

    return df, new_group

So that:

df, new_group = move_nan(df, 14.5)

print(df)
# Output
   group        numbers
0      0       [11, 12]
1      1  [14, nan, 15]
2      2       [16, 17]
3      3       [18, 19]

print(new_group)  # 1
user_112460560
Related