Extract a subset of a pandas dataframe based on values (with repetition)? - pandas

Question

Let say I have the following dataframe :
elements = [1,1,1,1,1,2,3,4,5]
df = pd.DataFrame({'elements': elements})
df.set_index(['elements'])
print df
elements
0 1
1 1
2 1
3 1
4 1
5 2
6 3
I have a list [1, 1, 2, 3] and I want a subset of the dataframe including those 4 elements, for example:
elements
0 1
1 1
5 2
6 3
I have been able to deal with it by building a dict counting the items occurrences in the array and building a new dataframe by appending subparts of the initial one.
Would you know some dataframe methods to help me find a more elegant solution?
After #jezrael comment : I must add that i need to keep track of the initial index (in df).
We can see df (first dataframe) as a repository of resources and i need to track which rows/indices are attributed :
Use case is : among the elements in df give me two 1, one 2 and one 3. i would persist the fact that i have the rows 0 and 1 as 1, row 4 as 2 and row 5 as 3.

Answer 1

If and only if your Series and list are sorted (otherwise, see below), then you can do:
L = [1, 1, 2, 3]
df[df.elements.apply(lambda x: x == L.pop(0) if x in L else False)]
elements
0 1
1 1
5 2
6 3
list.pop(i) returns and removes the value in list at index i. Because both, the elements and L, are sorted, popping the first element (i==0) of the subset list L will always occur at the corresponding first element in elements.
So at each iteration of lambda on elements, L will become:
| element | L | Output |
|=========|==============|===========|
| 1 | [1, 1, 2, 3] | True |
| 1 | [1, 2, 3] | True |
| 1 | [2, 3] | False |
| 1 | [2, 3] | False |
| 1 | [2, 3] | False |
| 2 | [2, 3] | True |
| 3 | [3] | True |
| 4 | [] | False |
| 5 | [] | False |
As you can see, your list is empty at the end, so if it's a problem, you can copy it beforehand. Or, you actually have that information in the new dataframe you just created!
If df.elements is not sorted, create a sorted copy on which you apply the same lambda function as above, but the output of it will be used as index for the original dataframe (indexes whose values are True are used):
df
elements
0 5
1 4
2 3
3 1
4 2
5 1
6 1
7 1
8 1
cp = df.elements.copy()
cp.sort_values(inplace=True)
tmp = df.loc[cp.apply(lambda x: x == L.pop(0) if x in L else False)]
print tmp
elements
2 3
3 1
4 2
5 1
HTH

Answer 2

Extracting can be possible by merge with new columns by GroupBy.cumcount:
L = [1,1,2,3]
df1 = pd.DataFrame({'elements':L})
df['g'] = df.groupby('elements')['elements'].cumcount()
df1['g'] = df1.groupby('elements')['elements'].cumcount()
print (df)
elements g
0 1 0
1 1 1
2 1 2
3 1 3
4 1 4
5 2 0
6 3 0
7 4 0
8 5 0
print (df1)
elements g
0 1 0
1 1 1
2 2 0
3 3 0
print (pd.merge(df,df1, on=['elements', 'g']))
elements g
0 1 0
1 1 1
2 2 0
3 3 0
print (pd.merge(df.reset_index(),df1, on=['elements', 'g'])
.drop('g', axis=1)
.set_index('index')
.rename_axis(None))
elements
0 1
1 1
5 2
6 3