How to split a column of items - which are not in order, into their own columns in pandas?

Viewed 67

So basically, I am looking at survey data, and one of the questions has an answer where people can select multiple options. All these options will now be under one column, but based on the order in which the respondent has selected the options. So basically, when I perform value_counts() on the column, it looks like this:

A       10
B       15
C       6
D       19
E       23
A,B     2
A,C     5
A,B,E   7
E,A,C   4
B,C     6
..

So now I want to select a combination where respondents have selected at least one of A,B,C or D to this question, but not just A, just B, just C or just D. So in essence, I want the combinations where the selected option is more than one, and the options have at least A/B/C/D. Ex: (A,D,E), (A,B,F), (B,F) and so on.

I have tried splitting this with a simple delimiter and making a column for each option, but the problem is: not all the rows are of same length and also, the order is not always the same and all the first elements go under the first column, which again makes it useless. I have tried manually selecting the options from value counts, like:

df = df[(df['variable'] == 'A,D,E') | (df['variable'] == 'B,F')]

But I want to be able to automate this, but just don't know how.

I have also tried looking for a sub-string, creating a column, and adding all the numbers - but the issue with this is, it also takes respondents who answered just A, just B, etc. into consideration, which is not something I want.

x ='A'
df["A_column"]= df["variable"].str.find(x)

Can someone please help me with the logic of this?

2 Answers

One solution would be to create a dictionary with each string of interest mapped to a character, and transform your variable column into the one encoded by characters. Then you can check whether any of the characters corresponding to the strings of interest are present and check whether the length of the transformed column is greater than one.

For example if you have the following df and the strings you want to detect are "str_a", "str_b", "str_c":

>>> df
                  variable  value
0                    str_a      0
1              str_a,str_b      1
2                    str_b      2
3        str_a,str_b,str_c      3
4                    str_c      4
5  str_a,str_b,str_c,str_d      5
6                    str_e      6
7              str_d,str_e      7
8              str_a,str_d      8
9        str_a,str_d,str_e      9

We can make the mapping:

variable_map = {'str_a': 'A', 'str_b': 'B', 'str_c': 'C', 'str_d': 'D', 'str_e': 'E'}

And transform your variable column:

df['variable_encoded'] = df['variable'].replace(variable_map, regex=True)

This will give us the following:

>>> df
                  variable  value variable_encoded
0                    str_a      0                A
1              str_a,str_b      1              A,B
2                    str_b      2                B
3        str_a,str_b,str_c      3            A,B,C
4                    str_c      4                C
5  str_a,str_b,str_c,str_d      5          A,B,C,D
6                    str_e      6                E
7              str_d,str_e      7              D,E
8              str_a,str_d      8              A,D
9        str_a,str_d,str_e      9            A,D,E

Then we can check for the presence of the values corresponding to the keys "str_a", "str_b", "str_c" which would be "A","B","C" and that the length of the row is also greater than 1 (to ensure we don't return any rows that only have one of the strings we want).

>>> df[(df['variable_encoded'].str.contains("A|B|C")) & (df['variable_encoded'].str.len() > 1)]

This gives us the following selection:

                  variable  value variable_encoded
1              str_a,str_b      1              A,B
3        str_a,str_b,str_c      3            A,B,C
5  str_a,str_b,str_c,str_d      5          A,B,C,D
8              str_a,str_d      8              A,D
9        str_a,str_d,str_e      9            A,D,E

(first I'm creating a dataframe I think is similar to yours)

df = pd.DataFrame([{'answer': 'A,B,C'}, {'answer': 'A,D'}, {'answer': 'B,C,D'}, {'answer': 'D,E'}, {'answer': 'E'}, {'answer': 'B,C'}])
df
  answer
0  A,B,C
1    A,D
2  B,C,D
3    D,E
4      E
5    B,C

first step is to turn your csv strings into an explodable column type

df2 = df.reset_index().rename(
    columns={'index': 'original_index'}
).pipe(lambda x: x.assign(
    split_answer=x.answer.apply(lambda a: a.split(',')),
    selected=True
))
df2
   original_index answer split_answer  selected
0               0  A,B,C    [A, B, C]      True
1               1    A,D       [A, D]      True
2               2  B,C,D    [B, C, D]      True
3               3    D,E       [D, E]      True
4               4      E          [E]      True
5               5    B,C       [B, C]      True

second you can explode the new column into multiple rows per answer and pivot the options out into columns

df2.explode('split_answer').pivot(columns='split_answer', index='original_index', values='selected').fillna(False)
split_answer        A      B      C      D      E
original_index                                   
0                True   True   True  False  False
1                True  False  False   True  False
2               False   True   True   True  False
3               False  False  False   True   True
4               False  False  False  False   True
5               False   True   True  False  False

the selected=True and .fillna(False) is just to get the final result to look nice with booleans

Related