adding a legend to animated scatterplot

Question

I am building an animation with a scatter plot showing data for multiple groups over time.
When I want to add a legend, the best I could get only shows one group.
sample dataset:

import pandas as pd 

df = pd.DataFrame([
    [1, 'a', 0.39, 0.73],
    [1, 'b', 0.87, 0.94],
    [1, 'c', 0.87, 0.23],
    [2, 'a', 0.17, 0.37],
    [2, 'b', 0.03, 0.12],
    [2, 'c', 0.86, 0.22],
    [3, 'a', 0.01, 0.15],
    [3, 'b', 0.03, 0.1],
    [3, 'c', 0.29, 0.19],
    columns=['period', 'group', 'x', 'y']
)

I build my animation like this:

import matplotlib.pyplot as plt
import matplotlib.animation as animation

fig, ax = plt.subplots()

ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
colors = {
'a': 'r',
'b': 'b',
'c': 'g'
        }
scat = ax.scatter([], [],
                c=df['group'].map(colors),
                )

def init():
    scat.set_offsets([])
    return scat,

def update(period):
    scat.set_offsets(df[df['period'] == period][['x', 'y']])
    scat.set_label(df[df['period']  == period]['group'])
    ax.legend([scat], df['group'].unique().tolist(), loc=1)
    ax.set_title(period)
    return scat,

ani = animation.FuncAnimation(fig, update, init_func=init,
                            frames=[1,2,3,4,5],
                            interval=500,
                            repeat=True)

plt.show()

I get only group a showing up in the legend.

If I only type ax.legend(loc=1), It shows something that looks like this:

6  a
7  b
8  c
Name: group, dtype:object

The numbers change in each frame.

I already checked those answers:
How do I get this to show the legend on the plot?: got me where I am now.
How to add legend/label in python animation: I get UnboundLocalError: local variable 'legend' referenced before assignment on legend.remove()
Add a legend for an animation (of Artists) in matplotlib: only shows group a.

Answer 1

I've found the solution.
I need to create one scatter plot per group. Then I update every scatter plot in my update() method.
Here's my final code:

fig, ax = plt.subplots()
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
colors = {
    'a': 'r',
    'b': 'b',
    'c': 'g'
}

scats = []
groups = df.groupby('group')
for name, grp in groups:
    scat = ax.scatter([], [],
                      color=colors[name],
                      label=name)
    scats.append(scat)
ax.legend(loc=4)

def init():
    for scat in scats:
        scat.set_offsets([])
    return scats,

def update(period):
    for scat, (name, data) in zip(scats, groups):
        sample = data[data['period'] == period][['x', 'y']]
        scat.set_offsets(sample)
return scats,

ani = animation.FuncAnimation(fig, update, init_func=init
                              frames=[1, 2, 3, 4, 5],
                              interval=500,
                              repeat=True)

plt.show()