python - creating a matplotlib scatter legend size related

Question

I am looking for a way to include a (matplotlib) legend that describe the size of points in a scatter plot, as this could be related to another variable, like in this basic example:

import numpy as np
import matplotlib.pyplot as plt

N = 50
x = np.random.rand(N)
y = np.random.rand(N)
a2 = 400*np.random.rand(N)

plt.scatter(x, y, s=a2, alpha=0.5)
plt.show()

(inspired from: http://matplotlib.org/examples/shapes_and_collections/scatter_demo.html)

so in the legend there would be ideally few spots corresponding to sizes 0-400 (the a2 variable), according to s descriptor in scatter.

Answer 1

Use .legend_elements("sizes"):

import numpy as np
import matplotlib.pyplot as plt

N = 50
x = np.random.rand(N)
y = np.random.rand(N)
a2 = 400*np.random.rand(N)

sc = plt.scatter(x, y, s=a2, alpha=0.5)
plt.legend(*sc.legend_elements("sizes", num=6))
plt.show()