CODEWITHSUNDEEP
pythonpandasdataframe
gal leshem
gal leshem

Reputation: 561

add column to groups on dataframe pandas

I have a dataframe:

  id  |   x   |   y 
   1  |  0.3  |  0.4
   1  |  0.2  |  0.5
   2  |  0.1  |  0.6
   2  |  0.9  |  0.1
   3  |  0.8  |  0.2
   3  |  0.7  |  0.3

How can I add a new column to dataframe relative to the id column?

for example: 

  id  |   x   |   y   |  color
   1  |  0.3  |  0.4  | 'green'
   1  |  0.2  |  0.5  | 'green'
   2  |  0.1  |  0.6  | 'black'
   2  |  0.9  |  0.1  | 'black'
   3  |  0.8  |  0.2  |  'red'
   3  |  0.7  |  0.3  |  'red'

Upvotes: 0

Views: 38

Answers (2)

Gius
Gius

Reputation: 514

Probably to late but if you want alternatives, here's another way with a simple function:

colors = ['Green', 'Black', 'Red']

def color(data):
    if data['id'] == 1:
        col = colors[0]
    if data['id'] == 2:
        col = colors[1]
    if data['id'] == 3:
        col = colors[2]
    return col

df['Colors'] = df.apply(color, axis = 1)
print(df)

#    id    x    y Colors
# 0   1  0.3  0.4  Green
# 1   1  0.2  0.5  Green
# 2   2  0.1  0.6  Black
# 3   2  0.9  0.1  Black
# 4   3  0.8  0.2    Red
# 5   3  0.7  0.3    Red

Upvotes: 0

Rafal Janik
Rafal Janik

Reputation: 309

So your function doesn't return color names but instead the RGB values, if this it what you want in the color column build the dictionary first from the unique id values and apply the dictionary the way @anky_91 mentioned in the comments.

d={x:random_color() for x in df.id.unique()}
df['color']=df['id'].map(d)

Upvotes: 1

Related Questions