在python中使用pandas根据其他列中给出的值选择列

Question

I have a data frame as:

a   b   c   d......

1   1
3   3   3   5
4   1   1   4   6
1   0

I want to select number of columns based on value given in column "a". In this case for first row it would only select column b. How can I achieve something like:

df.iloc[:,column b:number of columns corresponding to value in column a]

My expected output would be:

a   b   c   d   e
1   1   0   0   1     # 'e' contains value in column b because colmn a = 1 
3   3   3   5   335   #  'e' contains values of column b,c,d because colm a 
4   1   1   4   1      #  = 3
1   0           NAN

Answer 1

Define a little function for this:

def select(df, r):
    return df.iloc[r, 1:1 + df.iat[r, 0]]  

The function uses iat to query the a column for that row, and iloc to select columns from the same row.

Call it as such:

select(df, 0)

b    1.0
Name: 0, dtype: float64

And,

select(df, 1)

b    3.0
c    3.0
d    5.0
Name: 1, dtype: float64

---

Based on your edit, consider this -

df

   a  b  c  d  e
0  1  1  0  0  0
1  3  3  3  5  0
2  4  1  1  4  6
3  1  0  0  0  0

Use where / mask (with numpy broadcasting) + agg here -

df['e'] = df.iloc[:, 1:]\
            .astype(str)\
            .where(np.arange(df.shape[1] - 1) < df.a[:, None], '')\
            .agg(''.join, axis=1)

df

   a  b  c  d     e
0  1  1  0  0     1
1  3  3  3  5   335
2  4  1  1  4  1146
3  1  0  0  0     0

If nothing matches, then those entries in e will have an empty string. Just use replace -

df['e'] = df['e'].replace('', np.nan)

Answer 2

A numpy slicing approach

a = v[:, 0]
b = v[:, 1:]
n, m = b.shape
b = b.ravel()
b = np.where(b == 0, '', b.astype(str))
r = np.arange(n) * m
f = lambda t: b[t[0]:t[1]]

df.assign(g=list(map(''.join, map(f, zip(r, r + a)))))

   a  b  c  d  e     g
0  1  1  0  0  0     1
1  3  3  3  5  0   335
2  4  1  1  4  6  1146
3  1  0  0  0  0      

Answer 3

Edit : one line solution with slicing.

df["f"] = df.astype(str).apply(lambda r: "".join(r[1:int(r["a"])+1]), axis=1)

# df["f"] = df["f"].astype(int)  if you need `f` to be integer

df    
    a   b   c   d   e   f
0   1   1   X   X   X   1
1   3   3   3   5   X   335
2   4   1   1   4   6   1146
3   1   0   X   X   X   0

Dataset used:

df = pd.DataFrame({'a': {0: 1, 1: 3, 2: 4, 3: 1},
                   'b': {0: 1, 1: 3, 2: 1, 3: 0},
                   'c': {0: 'X', 1: '3', 2: '1', 3: 'X'},
                   'd': {0: 'X', 1: '5', 2: '4', 3: 'X'},
                   'e': {0: 'X', 1: 'X', 2: '6', 3: 'X'}})

Suggestion for improvement would be appreciated!