title: How to efficiently partial argsort Pandas dataframe across columns
body: <p>I would like to replace values with column labels according to the largest 3 values for each row. Let's assume this input:</p>
  
  <pre><code>   p1  p2  p3  p4
  0   0   9   1   4
  1   0   2   3   4
  2   1   3  10   7
  3   1   5   3   1
  4   2   3   7  10
  </code></pre>
  
  <p>Given <code>n = 3</code>, I am looking for:</p>
  
  <pre><code>  Top1 Top2 Top3
  0   p2   p4   p3
  1   p4   p3   p2
  2   p3   p4   p2
  3   p2   p3   p1
  4   p4   p3   p2
  </code></pre>
  
  <p>I'm not concerned about duplicates, e.g. for index <code>3</code>, <code>Top3</code> can be <code>'p1'</code> <em>or</em> <code>'p4'</code>.</p>
  
  <h3>Attempt 1</h3>
  
  <p>My first attempt is a full sort using <a href="https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.argsort.html" rel="noreferrer"><code>np.ndarray.argsort</code></a>:</p>
  
  <pre><code>res = pd.DataFrame(df.columns[df.values.argsort(1)]).iloc[:, len(df.index): 0: -1]
  </code></pre>
  
  <p>But in reality I have more than 4 columns and this will be inefficient.</p>
  
  <h3>Attempt 2</h3>
  
  <p>Next I tried <a href="https://docs.scipy.org/doc/numpy/reference/generated/numpy.argpartition.html" rel="noreferrer"><code>np.argpartition</code></a>. But since values within each partition are not sorted, this required a <em>subsequent</em> sort:</p>
  
  <pre><code>n = 3
  parts = np.argpartition(-df.values, n, axis=1)[:, :-1]
  args = (-df.values[np.arange(df.shape[0])[:, None], parts]).argsort(1)
  
  res = pd.DataFrame(df.columns[parts[np.arange(df.shape[0])[:, None], args]],
                     columns=[f'Top{i}' for i in range(1, n+1)])
  </code></pre>
  
  <p>This, in fact, works out <em>slower</em> than the first attempt for larger dataframes. Is there a more efficient way which takes advantage of partial sorting? You can use the below code for benchmarking purposes.</p>
  
  <h3>Benchmarking</h3>
  
  <pre><code># Python 3.6.0, NumPy 1.11.3, Pandas 0.19.2
  
  import pandas as pd, numpy as np
  
  df = pd.DataFrame({'p1': [0, 0, 1, 1, 2],
                     'p2': [9, 2, 3, 5, 3],
                     'p3': [1, 3, 10, 3, 7],
                     'p4': [4, 4, 7, 1, 10]})
  
  def full_sort(df):
      return pd.DataFrame(df.columns[df.values.argsort(1)]).iloc[:, len(df.index): 0: -1]
  
  def partial_sort(df):
      n = 3
      parts = np.argpartition(-df.values, n, axis=1)[:, :-1]
      args = (-df.values[np.arange(df.shape[0])[:, None], parts]).argsort(1)
      return pd.DataFrame(df.columns[parts[np.arange(df.shape[0])[:, None], args]])
  
  df = pd.concat([df]*10**5)
  
  %timeit full_sort(df)     # 86.3 ms per loop
  %timeit partial_sort(df)  # 158 ms per loop
  </code></pre>