42 lines
1.6 KiB
Python
42 lines
1.6 KiB
Python
import pandas as pd
|
||
from sklearn.preprocessing import LabelEncoder
|
||
from sklearn.manifold import TSNE
|
||
from sklearn.cluster import KMeans
|
||
import matplotlib.pyplot as plt
|
||
import seaborn as sns
|
||
|
||
# Загрузка данных
|
||
data = pd.read_csv("ds_salaries.csv")
|
||
|
||
# Преобразование категориальных переменных в числовые
|
||
le = LabelEncoder()
|
||
data['company_location'] = le.fit_transform(data['company_location'])
|
||
data['company_size'] = le.fit_transform(data['company_size'])
|
||
|
||
# Выбор нужных признаков
|
||
features = ['company_location', 'company_size']
|
||
|
||
# Применение t-SNE для уменьшения размерности
|
||
tsne = TSNE(n_components=2, random_state=42)
|
||
data_tsne = tsne.fit_transform(data[features])
|
||
|
||
# Кластеризация данных
|
||
kmeans = KMeans(n_clusters=3, random_state=42)
|
||
labels = kmeans.fit_predict(data_tsne)
|
||
|
||
# Создание датафрейма с новыми координатами и метками кластеров
|
||
data_tsne_df = pd.DataFrame(data=data_tsne, columns=['Dimension 1', 'Dimension 2'])
|
||
data_tsne_df['Cluster'] = labels
|
||
|
||
# Добавление номера кластера в исходные данные
|
||
data['Cluster'] = labels
|
||
|
||
# Визуализация
|
||
plt.figure(figsize=(10, 6))
|
||
sns.scatterplot(x='Dimension 1', y='Dimension 2', hue='Cluster', data=data_tsne_df, palette='viridis', s=50)
|
||
plt.title('t-SNE Clustering of Companies')
|
||
plt.show()
|
||
|
||
# Вывод номера кластера в исходных данных
|
||
print(data[['company_location', 'company_size', 'Cluster']])
|