Submission #80861

ソースコード

?
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
#include<bits/stdc++.h>
using namespace std;
struct SparseTable{
vector<int>a;
vector<vector<int>>table;
vector<int>log_table;
SparseTable(){};
void build(vector<int> &v){
a = v;
int log_n = 0;
int siz = v.size();
while((1<<log_n) <= siz)++log_n;
int n = 1 << log_n;
table.resize(n,vector<int>(log_n));
a.resize(n);
log_table.resize(siz + 1);
for(int i = 0;i < siz;i++){
table[i][0] = i;
}
for(int k = 1;k < log_n;k++){
for(int i = 0;i+(1<<k) <= n;i++){
if(a[table[i][k - 1]] < a[table[i + (1 << k - 1)][k - 1]]){
table[i][k] = table[i][k - 1];
}else{
table[i][k] = table[i + (1 << k - 1)][k - 1];
}
}
}
for(int i = 2; i < siz + 1; ++i){
log_table[i] = log_table[i >> 1] + 1;
}
}
//半開区間
int query_idx(int l,int r){
int log_n = log_table[r - l];
if(a[table[l][log_n]] < a[table[r - (1 << log_n)][log_n]]){
return (table[l][log_n]);
}else{
return (table[r - (1 << log_n)][log_n]);
}
}
};
struct LCA{
using Graph = vector<vector<int>>;
int n;
Graph graph;
vector<int>depth,index,firstpos;
SparseTable st;
LCA(Graph G){
graph = G;
n = G.size();
firstpos.resize(n,INT_MAX);
build();
}
void build(int s = 0){
dfs(s,-1,0);
st.build(depth);
return;
}
void dfs(int now,int par,int dep){
firstpos[now] = min(firstpos[now],(int)index.size());
depth.push_back(dep);
index.push_back(now);
for(auto i:graph[now]){
if(i != par){
dfs(i,now,dep+1);
depth.push_back(dep);
index.push_back(now);
}
}
}
int query(int u,int v){
u = firstpos[u];
v = firstpos[v];
if(u > v)swap(u,v);
return index[st.query_idx(u,v+1)];
return 0;
}
int operator[](int &p){
return depth[firstpos[p]];
}
};
//AuxiliaryTreeのコード
struct AuxiliaryTree{
using Graph = vector<vector<int>>;
Graph graph;
LCA lca;
int n,cnt;
vector<int>pre_order;
vector<int>vec;
Graph aux;
AuxiliaryTree(Graph &G):lca(G),graph(G){
n = graph.size();
pre_order.resize(n);
aux.resize(n);
cnt = 0;
dfs(0,-1);
}
void dfs(int now,int par){
pre_order[now] = cnt;
++cnt;
for(auto i:graph[now]){
if(par != i){
dfs(i,now);
}
}
}
void query(vector<int>&v){
int siz = v.size();
sort(v.begin(),v.end(),[&](int a,int b){return (pre_order[a] < pre_order[b]);});
for(int i = 0;i < siz-1;i++){
v.push_back(lca.query(v[i],v[i+1]));
}
sort(v.begin(),v.end(),[&](int a,int b){return (pre_order[a] < pre_order[b]);});
v.erase(unique(v.begin(),v.end()),v.end());
siz = v.size();
for(int i = 0;i < siz;i++){
aux[v[i]].clear();
}
for(int i = 0;i < siz-1;i++){
int x = lca.query(v[i],v[i+1]);
if(x == v[i]){
aux[x].push_back(v[i+1]);
aux[v[i+1]].push_back(x);
}else if(x == v[i+1]){
aux[x].push_back(v[i]);
aux[v[i]].push_back(x);
}else{
aux[x].push_back(v[i+1]);
aux[v[i+1]].push_back(x);
}
}
return;
}
};
int main(){
cin.tie(nullptr);
ios_base::sync_with_stdio(false);
cout.setf(ios_base::fixed);
cout.precision(10);
int n;
cin >>n;
vector<int>c(n);
vector<vector<int>>graph(n);
vector<vector<int>>colors(n);
for(int i = 0;i < n;i++){
cin >>c[i];
c[i]--;
colors[c[i]].push_back(i);
}
for(int i = 0;i < n-1;i++){
int s,t;
cin >>s>>t;
s--,t--;
graph[s].push_back(t);
graph[t].push_back(s);
}
AuxiliaryTree aux(graph);
vector<int>dist(n),pre(n);
vector<int>res(n,INT_MAX);
for(int i = 0;i < n;i++){
if(colors[i].empty())continue;
vector<int>node = colors[i];
aux.query(node);
priority_queue<pair<int, int>, vector<pair<int, int> >, greater<pair<int, int> > > pq;
for(auto j:node){
dist[j] = INT_MAX;
pre[j] = -1;
}
for(auto j:colors[i]){
dist[j] = 0;
pre[j] = j;
pq.push({0,j});
}
while(!pq.empty()){
auto [now_cost,now_node] = pq.top();
pq.pop();
if(dist[now_node] < now_cost)continue;
for(auto next_node:aux.aux[now_node]){
int dis = abs(aux.lca[now_node] - aux.lca[next_node]);
if(dist[now_node] + dis < dist[next_node]){
dist[next_node] = dist[now_node] + dis;
pre[next_node] = pre[now_node];
pq.push({dist[next_node],next_node});
}
}
}
for(auto u:node){
for(auto v:aux.aux[u]){
u = pre[u];
v = pre[v];
if(u == v)continue;
int lca = aux.lca.query(u,v);
int dis = aux.lca[u] + aux.lca[v] - aux.lca[lca] * 2;
if(dis < res[u]){
res[u] = dis;
}
if(dis < res[v]){
res[v] = dis;
}
}
}
}
for(auto i:res){
cout <<i<<"\n";
}
}

ステータス

項目 データ
問題 1448 - 色付き山小屋
ユーザー名 ei2332
投稿日時 2024-08-29 20:31:20
言語 C++17
状態 Accepted
得点 13
ソースコード長 4920 Byte
最大実行時間 544 ms
最大メモリ使用量 139608 KB

セット

セット 得点 Cases
1 ALL 13 / 13 *

テストケース

ファイル名 状態 実行時間 メモリ使用量 #
in1 AC 17 ms 604 KB
1
in2 AC 18 ms 620 KB
1
in3 AC 24 ms 7336 KB
1
in4 AC 20 ms 16036 KB
1
in5 AC 19 ms 16672 KB
1
in6 AC 21 ms 16756 KB
1
in7 AC 22 ms 16712 KB
1
in8 AC 18 ms 16668 KB
1
in9 AC 24 ms 16616 KB
1
in10 AC 15 ms 688 KB
1
in11 AC 24 ms 624 KB
1
in12 AC 23 ms 436 KB
1
in13 AC 23 ms 500 KB
1
in14 AC 20 ms 556 KB
1
in15 AC 32 ms 996 KB
1
in16 AC 31 ms 1096 KB
1
in17 AC 24 ms 1068 KB
1
in18 AC 18 ms 1040 KB
1
in19 AC 40 ms 7044 KB
1
in20 AC 34 ms 7116 KB
1
in21 AC 219 ms 58020 KB
1
in22 AC 226 ms 58880 KB
1
in23 AC 537 ms 124044 KB
1
in24 AC 490 ms 125044 KB
1
in25 AC 457 ms 127980 KB
1
in26 AC 455 ms 132092 KB
1
in27 AC 455 ms 134516 KB
1
in28 AC 478 ms 134608 KB
1
in29 AC 450 ms 132488 KB
1
in30 AC 471 ms 138440 KB
1
in31 AC 425 ms 138944 KB
1
in32 AC 425 ms 136208 KB
1
in33 AC 404 ms 133532 KB
1
in34 AC 398 ms 132100 KB
1
in35 AC 450 ms 132968 KB
1
in36 AC 465 ms 139608 KB
1
in37 AC 454 ms 138368 KB
1
in38 AC 462 ms 138296 KB
1
in39 AC 470 ms 138348 KB
1
in40 AC 544 ms 138268 KB
1