CODEWITHSUNDEEP
pythonpandasdataframe
jonas
jonas

Reputation: 13969

Create ID column in a pandas dataframe

I have a dataframe containing a trading log. My problem is that I do not have any ID to match buy and sell of a stock. The stock could be traded many times and I would like to have an ID to match each finished trade. My original dataframe a sequential timeseries dataframe with timestamps. The below example illustrates my problem, I need to match and ID traded stock in sequential order. Very simplified example:

df1 = pd.DataFrame({'stock': ['A', 'B', 'C', 'A','C', 'A', 'A'],
                        'deal': ['buy', 'buy', 'buy', 'sell','sell', 'buy', 'sell']}) 
df1
Out[84]: 
  stock  deal
0     A   buy
1     B   buy
2     C   buy
3     A  sell
4     C  sell
5     A   buy
6     A  sell

Here is my desired output:

df1 = pd.DataFrame({'stock': ['A', 'B', 'C', 'A','C', 'A', 'A'],
                    'deal': ['buy', 'buy', 'buy', 'sell','sell', 'buy', 'sell'],
                    'ID': [1, 2, 3, 1,3, 4, 4]}) 


df1
Out[82]: 
  stock  deal  ID
0     A   buy   1
1     B   buy   2
2     C   buy   3
3     A  sell   1
4     C  sell   3
5     A   buy   4
6     A  sell   4

Any ideas?

Upvotes: 5

Views: 5251

Answers (2)

Mahir Mahbub
Mahir Mahbub

Reputation: 135

Try This:

import pandas as pd
df1 = pd.DataFrame({'stock': ['A', 'B', 'C', 'A','C', 'A', 'A'],
                'deal': ['buy', 'buy', 'buy', 'sell','sell', 'buy', 'sell']})

def sequential_buy_sell_id_generator(df1):

    column_length = len(df1["stock"])
    found = [0]*column_length
    id = [0]*column_length

    counter = 0

    for row_pointer_head in range(column_length):
        if df1["deal"][row_pointer_head]=="buy":
            id[row_pointer_head]= counter
            counter+=1
            found[row_pointer_head] = 1
            id[row_pointer_head]= counter

            for row_pointer_tail in range(row_pointer_head+1, column_length):

                if df1["stock"][row_pointer_head]== df1["stock"][row_pointer_tail] and df1["deal"][row_pointer_tail] =="sell" and found[row_pointer_tail] == 0:
                    found[row_pointer_tail] = 1
                    id[row_pointer_tail]= counter
                    break

    df1 = df1.assign(id = id) 
    return df1


print(sequential_buy_sell_id_generator(df1))

Output:

enter code here
    stock  deal  id
0     A   buy   1
1     B   buy   2
2     C   buy   3
3     A  sell   1
4     C  sell   3
5     A   buy   4
6     A  sell   4

Another Example:

For df1 = pd.DataFrame({'stock': ['A', 'B', 'C', 'A','C', 'A', 'A'],
                'deal': ['buy', 'buy', 'buy', 'buy','sell', 'sell', 'sell']})
  stock deal    ID
0   A   buy     1
1   B   buy     2
2   C   buy     3
3   A   buy     4
4   C   sell    3
5   A   sell    1
6   A   sell    4

Upvotes: 1

Scott Boston
Scott Boston

Reputation: 153460

Try this:

m = df1['deal'] == 'buy'
df1['ID'] = m.cumsum().where(m)
df1['ID'] = df1.groupby('stock')['ID'].ffill()

df1

Output:

  stock  deal   ID
0     A   buy  1.0
1     B   buy  2.0
2     C   buy  3.0
3     A  sell  1.0
4     C  sell  3.0
5     A   buy  4.0
6     A  sell  4.0

Details:

  • Create a boolean series, True where deal equals 'buy'
  • Cumsum and assign to 'ID' to buy records
  • Use groupby and ffill to assign 'ID' to next 'sell' record buy 'stock'

Upvotes: 3

Related Questions