CODEWITHSUNDEEP
pythonlistconditional-statements
user3397243
user3397243

Reputation: 567

Python, appending to a list using conditions

I have a file with two columns, lets say A and B

A     B
1     10
0     11
0     12 
0     15
1     90
0     41

I want to create a new column (a list), lets call the empty list C = []

I would like to loop through A, find if A == 1, and if it is I want to append the value of B[A==1] (10 in the first case) to C until the next A == 1 arrives.

So my final result would be:

A     B     C
1     10    10
0     11    10
0     12    10
0     15    10
1     90    90
0     41    90

I have tried using the for loop, but only to my dismay:

for a in A:
    if a == 1:
        C.append(B[a==1])
    elif a == 0:
        C.append(B[a==1])

Upvotes: 0

Views: 3711

Answers (2)

ovgolovin
ovgolovin

Reputation: 13420

You may also try using groupby.

Though solution I have come up with looks a bit convoluted to me:

>>> from itertools import izip, groupby, count
>>> from operator import itemgetter
>>> def gen_group(L):
    acc = 0
    for item in L:
        acc += item
        yield acc


>>> [number_out for number,length in ((next(items)[1], 1 + sum(1 for _ in items)) for group,items in groupby(izip(gen_group(A), B), itemgetter(0))) for number_out in repeat(number, length)]
[10, 10, 10, 10, 90, 90]

The idea is to prepare groups and then use them to group your input:

>>> list(gen_group(A))
[1, 1, 1, 1, 2, 2]

Upvotes: 1

OzTamir
OzTamir

Reputation: 131

You could use another variable to keep the value of the last index in A that had a value of 1, and update it when the condition is met:

temp = 0
for index, value in enumerate(A):
    if value == 1:
        C.append(B[index])
        temp = index
    else:
        C.append(B[temp])

enumerate() gives you a list of tuples with index and values from an utterable. For A, it will be [(0, 1), (1, 0), (2, 0), (3, 0), (4, 1), (5, 0)].

P.S: When you try to address a list using a boolean (B[a == 1]) it will return the item in the first place if the condition is false (B[a != 1] => B[False] => B[0]) or the item in the second place if it's true (B[a == 1] => B[True] => B[1]).

Upvotes: 4

Related Questions