tz_atcoder/AHC020.py at main · tztechno/tz_atcoder · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
####################################
## RE,grid search version

#library block
from sklearn import svm, datasets
from sklearn.model_selection import GridSearchCV
import sys
input = sys.stdin.readline

n,m,k=map(int,input().split())
XY=[]
for i in range(n):
    x,y=map(int,input().split())
    XY+=[(x,y)]
UVW=[]
for i in range(m):
    u,v,w=map(int,input().split())
    UVW+=[(u,v,w)]
AB=[]
for i in range(k):
    a,b=map(int,input().split())
    AB+=[(a,b)]

# Define the objective function
def objective(parameters):
    MS = [1] * m
    PS = [10] * n
    cost = 0
    for i in range(m):
        MS[i] = parameters['m' + str(i)]
        cost += UVW[i][2] * MS[i]
    for j in range(n):
        PS[j] = parameters['p' + str(j)]
        cost += PS[j] ** 2
    return cost

# Define the parameter grid for GridSearchCV
parameter_grid = {}
for i in range(m):
    parameter_grid['m' + str(i)] = [0, 1]
for j in range(n):
    parameter_grid['p' + str(j)] = list(range(1000))

# Perform grid search
grid_search = GridSearchCV(estimator=svm.SVC(),
                           param_grid=parameter_grid,
                           scoring='neg_mean_squared_error',
                           cv=5)

grid_search.fit(XY, AB)
best_params = grid_search.best_params_

P = []
for i in range(n):
    P.append(best_params['p' + str(i)])
B = []
for i in range(m):
    B.append(best_params['m' + str(i)])

print(*P)
print(*B)

####################################
## RE, optuna version

import sys
input = sys.stdin.readline
import optuna

n,m,k=map(int,input().split())

XY=[]
for i in range(n):
    x,y=map(int,input().split())
    XY+=[(x,y)]

UVW=[]
for i in range(m):
    u,v,w=map(int,input().split())
    UVW+=[(u,v,w)]

AB=[]
for i in range(k):
    a,b=map(int,input().split())
    AB+=[(a,b)]

def objective(trial):
    MS=[1]*m
    PS=[10]*n
    cost=0
    for i in range(m):
        MS[i]=trial.suggest_categorical('m'+str(i),[0,1])
        cost+=UVW[i][2]*MS[i]
    for j in range(n):
        PS[i]=trial.suggest_int('p'+str(j),0,10000)
        cost+=PS[i]**2
    return cost

study = optuna.create_study(direction='minimize')
study.optimize(objective, n_trials=100)
best_trial= study.best_trial.params

P=[]
for i in range(n):
    P+=[best_trial['p'+str(i)]]
B=[]
for i in range(m):
    B+=[best_trial['m'+str(i)]]
print(*P)
print(*B)
####################################