CODEWITHSUNDEEP
pythonpandasdata-science
Grigorij Abramov
Grigorij Abramov

Reputation: 161

Select Pandas Columns with a List

df = pd.DataFrame({'a':np.random.randn(1000), 'b':np.random.randn(1000), 'c':np.random.randn(1000), 'd':np.random.randn(1000),
                  'e':np.random.randn(1000), 'f':np.random.randn(1000), 'g':np.random.randn(1000), 'h':np.random.randn(1000),
                  'i':np.random.randn(1000), 'j':np.random.randn(1000), 'k':np.random.randn(1000), 'l':np.random.randn(1000)})

I have a Dataframe like this with (a lot of) columns.

col_of_interest = ['a','f', 'j', 'k', 'c']

I want to choose only a subset which is saved in a list.

%timeit df.loc[:, df.columns.isin(col_of_interest)]
%timeit df[col_of_interest]
%timeit df[[c for c in df.columns if c in col_of_interest]]
%timeit df[np.intersect1d(df.columns, col_of_interest)]
%timeit df[df.columns & col_of_interest]

803 µs ± 289 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.92 ms ± 324 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
2.18 ms ± 406 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
2.53 ms ± 194 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.39 ms ± 34.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

I was a little surprised that df[col_of_interest] is not the best solution and was interested if there is an even better/efficient way than df.loc[:, df.columns.isin(col_of_interest)].

Upvotes: 1

Views: 81

Answers (1)

yatu
yatu

Reputation: 88236

On my machine np.in1d runs slightly faster, although the differences are negligible for relatively small amounts of columns:

%timeit df.loc[:, np.in1d(df.columns, col_of_interest)]
# 493 µs ± 2.59 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

You may also check the performance using sets, which will presumably reduce the time complexity for larger amounts of columns:

%timeit df.loc[:, set(df.columns).intersection(col_of_interest)]
# 915 µs ± 65.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Upvotes: 2

Related Questions