matplotlibでscatterの点をクリックすると色を変えたい

初めに

scatterでプロットしたところまでは良かったのですが、その後データの更新しようとすると難しかったので共有しておきます。

コード

scatter.py

import matplotlib.pyplot as plt
import itertools #デカルト積
import numpy as np

if __name__ == '__main__':
    x_max = 10
    y_max = 10

    fig = plt.figure(figsize=(5,5))
    ax = fig.add_subplot(111)

    ax.set_xlim(-x_max*0.05,(x_max-1)*1.05)
    ax.set_ylim(-y_max*0.05,(y_max-1)*1.05)

    x_data = [i for i in range(x_max)]
    y_data = [i for i in range(y_max)]
    tmp_data = list(itertools.product(x_data,y_data))
    t_x = [i[0] for i in tmp_data]
    t_y = [i[1] for i in tmp_data]
    data = [t_x,t_y]
    
    colors = [(0,0,0,1) for i in range(x_max*y_max)]
    art = ax.scatter(data[0],data[1],c=colors)
    
    def onclick(event):
        x = round(event.xdata,0)
        y = round(event.ydata,0)
        tmp = int(x*(y_max) + y)
        if colors[tmp] == (0,0,0,1):
            colors[tmp] = (0.5,0.5,0.5,1)
        elif colors[tmp] == (0.5,0.5,0.5,1):
            colors[tmp] = (0,0,0,1)
        art.set_facecolor(colors)
        plt.draw()


    plt.connect('button_press_event', onclick)

    plt.show()