title: Return Row Mean Based on Max of Other Rows in Python
body: <p>I have this DataFrame</p>
  <pre><code>         A  B  C  A1  B1 C1 
  1/1/2021 1  2  7  9   5  7
  1/2/2021 4  3  5  3   4  5
  1/3/2021 4  6  4  6   7  2  
  </code></pre>
  <p>I want to add a new column D that for each row returns the mean two columns across A,B,C based on the max two values of A1, B1, C1. So if A1 and B1 are larger than C1 then column D would equal the mean of column A and B.</p>
  <p>Expected Output:</p>
  <pre><code>         A  B  C  A1  B1 C1 D
  1/1/2021 1  2  7  9   5  7  4 (mean of A and C, since A1 and C1 are the top two)
  1/2/2021 4  3  5  3   4  5  4 (mean of B and C, since B1 and C1 are the top two)
  1/3/2021 4  6  4  6   7  2  5 (mean of A and B, since A1 and B1 are the top two)
  </code></pre>
  <p>I think I can achieve the results using a function like the one below (I just included the first part as an example), writing out all the combinations. But I want something that can be used with a large number of columns that will change and ideally I could adjust the TopN number. For example, get the average of the top 3 or 4 instead of the top 2. The columns would always be structured consistently and in the correct order. For example, 5 columns of data to be used for the average values and 5 columns of data in the same order to be used to determine the max values.</p>
  <pre><code>def maxcol(row):  
      if row[A1] &gt;= row[B1] and row[A1] &gt;= row[C1] and row[B1] &gt;= row[C1]:
           val = row[A] + row[B] / 2
      elif: 
           etc etc.
      return val
  </code></pre>
  <p>Is there a simple way to accomplish this without the brute force approach above?</p>
  <p>UPDATE: I updated the answer to a more general code that works for multiple columns and multiple numbers of top columns.</p>
  <pre><code>import heapq
  
  df = pd.DataFrame({'A': [1, 4, 4],'B': [2, 3, 6],'C': [7, 5, 4],'A1': [9, 3, 6],'B1': [5, 4, 7],'C1': [7, 5, 2]})
  
  n = 3
  t = 2
  
  def helper(row):
      lst = [col for col in row]
      order = [lst[n:].index(x) for x in lst[n:] if x in heapq.nlargest(t,lst[n:])]
      return mean(lst[o] for o in order)
  
  df['D'] = df.apply(helper, axis = 1)
  print(df)
  </code></pre>