Working with few elements in a list belonging to a pandas dataframe

user_9217863 0
Viewed 197

Consider that I already figure it out a way to solve my problem, but i would like to find a better way and/or understanding a bit more about pandas dataframe and apply/lambda functions. I have a pandas DataFrame like this:

                      A
0             [43,45,51,...,32]
1             [43,45,51,...,32]
2             [43,45,51,...,32]
...                  ...
...                  ...
...                  ...
n             [43,45,51,...,32]

The column A contains per each row 64 elements. What I would like is to apply a function to each row that:

  1. sum only certain elements of the column A (let's say indexes 18,27,34,45)
  2. make the average of that sum
  3. put the average in a new column called B

so that, at the end I should have something like this:

                      A                      B (mean over few elements)
0             [43,45,51,...,32]              50
1             [42,55,41,...,68]              60
2             [11,30,29,...,23]              30
...                  ...                    ...
...                  ...                    ...
...                  ...                    ...
n             [33,45,84,...,77]              70

Right now I was able to do that using a numpy array in which i stored column A like this:

a = np.array(df["A"])
k =[18,27,34,45]
for i in range(n):
    vector = 0
    for idx, val in enumerate(k):
        vector = vector + a[i][val]
    mean.append(vector/16)

But i was thinking about a more elegant way with the apply/lamdba method. It will be also ok if, instead of creating another column I will append values to a numpy list. Thank you in advance and cheers.

pandas apply dataframe lambda python
3 Answers

Create a np.array from the column A, then use indexing in the array with indices k to select the desired elements and use np.mean along axis=1 and assign back this result to column B:

a = np.array(df['A'].tolist())
df['B'] = a[:, k].mean(axis=1)

Example:

np.random.seed(123)
df = pd.DataFrame({'A': [np.random.randint(1, 10, 10) for _ in range(5)]})

# print(df)
                                A
0  [3, 3, 7, 2, 4, 7, 2, 1, 2, 1]
1  [1, 4, 5, 1, 1, 5, 2, 8, 4, 3]
2  [5, 8, 3, 5, 9, 1, 8, 4, 5, 7]
3  [2, 6, 7, 3, 2, 9, 4, 6, 1, 3]
4  [7, 3, 5, 5, 7, 4, 1, 7, 5, 8]


k =[1, 2, 4, 7]
a = np.array(df['A'].tolist()) 
df['B'] = a[:, k].mean(axis=1)

Result:

                                A     B
0  [3, 3, 7, 2, 4, 7, 2, 1, 2, 1]  3.75
1  [1, 4, 5, 1, 1, 5, 2, 8, 4, 3]  4.50
2  [5, 8, 3, 5, 9, 1, 8, 4, 5, 7]  6.00
3  [2, 6, 7, 3, 2, 9, 4, 6, 1, 3]  5.25
4  [7, 3, 5, 5, 7, 4, 1, 7, 5, 8]  5.50
user_128331660 
import pandas


def column_generator(row_value, indices):
    # get a sublist of just the values that appear in the indices
    sublist = [row_value[index] for index in indices]
    # return average 
    return sum(sublist) / len(sublist)


df['B'] = df['A'].apply(column_generator, indices=[0, 1])
user_104410020

Try this:

indx = [0,2]
func = [lambda x: np.mean([x[i] for i in indx])]
df["mean"] = df.apply(func)
user_139287520
Related