[BUG] plot_clusters TypeError

[rohrdr]

🐛 Bug Report

Function plot_clusters gives error TypeError: 'AxesSubplot' object is not subscriptable when number of clusters is two.

Expected behavior

Should plot without errors

How To Reproduce

def gen_dataset():
    df = pd.DataFrame()
    for i in range(1, 5):
        date_range = pd.date_range("2020-01-01", "2020-05-01")
        for j, sigma in enumerate([0.1, 0.3, 0.5, 0.8]):
            tmp = pd.DataFrame({"timestamp": date_range})
            tmp["segment"] = f"{2*i}{j}"
            tmp["target"] = np.random.normal(2*i, sigma, len(tmp))
            df = df.append(tmp, ignore_index=True)
    ts = TSDataset(df=TSDataset.to_dataset(df), freq="D")
    return ts

ts = gen_dataset()
model = EuclideanClustering()
model.build_distance_matrix(ts=ts)
model.build_clustering_algo(n_clusters=2,linkage="average")
segment2cluster = model.fit_predict()
centroids = model.get_centroids()
plot_clusters(ts,segment2cluster,centroids)

Environment

No response

Additional context

No response

Checklist

• Bug appears at the latest library version