Submission #468823

#	Time	Username	Problem	Language	Result	Execution time	Memory
468823		hoaphat1	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	776 ms	55340 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
 
using namespace std;
 
template <typename T>
class graph {
 public:
  struct edge {
    int from;
    int to;
    T cost;
  };
 
  vector<edge> edges;
  vector<vector<int>> g;
  int n;
 
  graph(int _n) : n(_n) {
    g.resize(n);
  }
 
  virtual int add(int from, int to, T cost) = 0;
};
 
template <typename T>
class forest : public graph<T> {
 public:
  using graph<T>::edges;
  using graph<T>::g;
  using graph<T>::n;
 
  forest(int _n) : graph<T>(_n) {
  }
 
  int add(int from, int to, T cost = 1) {
    assert(0 <= from && from < n && 0 <= to && to < n);
    int id = (int) edges.size();
    assert(id < n - 1);
    g[from].push_back(id);
    g[to].push_back(id);
    edges.push_back({from, to, cost});
    return id;
  }
};
 
template <typename T>
class dfs_forest : public forest<T> {
 public:
  using forest<T>::edges;
  using forest<T>::g;
  using forest<T>::n;
 
  vector<int> pv;
  vector<int> pe;
  vector<int> order;
  vector<int> pos;
  vector<int> end;
  vector<int> sz;
  vector<int> root;
  vector<int> depth;
  vector<T> dist;
 
  dfs_forest(int _n) : forest<T>(_n) {
  }
 
  void init() {
    pv = vector<int>(n, -1);
    pe = vector<int>(n, -1);
    order.clear();
    pos = vector<int>(n, -1);
    end = vector<int>(n, -1);
    sz = vector<int>(n, 0);
    root = vector<int>(n, -1);
    depth = vector<int>(n, -1);
    dist = vector<T>(n);
  }
 
  void clear() {
    pv.clear();
    pe.clear();
    order.clear();
    pos.clear();
    end.clear();
    sz.clear();
    root.clear();
    depth.clear();
    dist.clear();
  }
 
 private:
  void do_dfs(int v) {
    pos[v] = (int) order.size();
    order.push_back(v);
    sz[v] = 1;
    for (int id : g[v]) {
      if (id == pe[v]) {
        continue;
      }
      auto &e = edges[id];
      int to = e.from ^ e.to ^ v;
      depth[to] = depth[v] + 1;
      dist[to] = dist[v] + e.cost;
      pv[to] = v;
      pe[to] = id;
      root[to] = (root[v] != -1 ? root[v] : to);
      do_dfs(to);
      sz[v] += sz[to];
    }
    end[v] = (int) order.size() - 1;
  }
 
  void do_dfs_from(int v) {
    depth[v] = 0;
    dist[v] = T{};
    root[v] = v;
    pv[v] = pe[v] = -1;
    do_dfs(v);
  }
 
 public:
  void dfs(int v, bool clear_order = true) {
    if (pv.empty()) {
      init();
    } else {
      if (clear_order) {
        order.clear();
      }
    }
    do_dfs_from(v);
  }
 
  void dfs_all() {
    init();
    for (int v = 0; v < n; v++) {
      if (depth[v] == -1) {
        do_dfs_from(v);
      }
    }
    assert((int) order.size() == n);
  }
};
 
template <typename A, typename B>
string to_string(pair<A, B> p);
 
template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p);
 
template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p);
 
string to_string(const string& s) {
  return '"' + s + '"';
}
 
string to_string(const char* s) {
  return to_string((string) s);
}
 
string to_string(bool b) {
  return (b ? "true" : "false");
}
 
string to_string(vector<bool> v) {
  bool first = true;
  string res = "{";
  for (int i = 0; i < static_cast<int>(v.size()); i++) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(v[i]);
  }
  res += "}";
  return res;
}
 
template <size_t N>
string to_string(bitset<N> v) {
  string res = "";
  for (size_t i = 0; i < N; i++) {
    res += static_cast<char>('0' + v[i]);
  }
  return res;
}
 
template <typename A>
string to_string(A v) {
  bool first = true;
  string res = "{";
  for (const auto &x : v) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(x);
  }
  res += "}";
  return res;
}
 
template <typename A, typename B>
string to_string(pair<A, B> p) {
  return "(" + to_string(p.first) + ", " + to_string(p.second) + ")";
}
 
template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " + to_string(get<2>(p)) + ")";
}
 
template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " + to_string(get<2>(p)) + ", " + to_string(get<3>(p)) + ")";
}
 
void debug_out() { cerr << endl; }
 
template <typename Head, typename... Tail>
void debug_out(Head H, Tail... T) {
  cerr << " " << to_string(H);
  debug_out(T...);
}
 
#ifndef LOCAL
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#else
#define debug(...) 42
#endif
/// template from tourist
 
int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  int n, m;
  cin >> n >> m;
  dfs_forest<int> g(n);
  for (int i = 1; i < n; i++) {
  	int u, v;
  	cin >> u >> v;
  	--u; --v;
  	g.add(u, v);
	}
	g.dfs(0);
	vector<int> dia(2);
	dia[0] = max_element(g.depth.begin(), g.depth.end()) - g.depth.begin();
	g.dfs(dia[0]);
	dia[1] = max_element(g.depth.begin(), g.depth.end()) - g.depth.begin();
	vector<int> c(n);
	for (int i = 0; i < n; i++) {
		cin >> c[i];
		--c[i];
	}
	vector<int> ans(n);
	for (int i = 0; i < 2; i++) {
		vector<pair<int, int>> dis(n);
		function<void(int, int)> dfs = [&](int v, int p) {
			dis[v] = make_pair(0, v);
			for (int id : g.g[v]) {
				auto& e = g.edges[id];
				int u = e.from ^ e.to ^ v;
				if (u == p) continue;
				dfs(u, v);
				dis[v] = max(dis[v], make_pair(dis[u].first + 1, u));
			}
		};
		vector<int> st;
		vector<int> cnt(m);
		int res = 0;
		auto add = [&](int v) {
			if (!cnt[c[v]]++) ++res;
		};
		auto del = [&](int v) {
			if (!--cnt[c[v]]) --res;
		};
		function<void(int, int)> dfs1 = [&](int v, int p) {
			if (p != -1) {
				st.push_back(p);
				add(p);
			}
			int mx = 0;
			for (int id : g.g[v]) {
				auto& e = g.edges[id];
				int u = e.from ^ e.to ^ v;
				if (u == p || u == dis[v].second) continue;
				mx = max(mx, dis[u].first + 1);
			}
			while (!st.empty() && g.depth[v] - g.depth[st.back()] <= mx) {
				del(st.back());
				st.pop_back();
			}
			if (v != dis[v].second) {
				dfs1(dis[v].second, v);
			}
			while (!st.empty() && g.depth[v] - g.depth[st.back()] <= dis[v].first) {
				del(st.back());
				st.pop_back();
			}
			for (int id : g.g[v]) {
				auto& e = g.edges[id];
				int u = e.from ^ e.to ^ v;
				if (u == p || u == dis[v].second) continue;
				dfs1(u, v);
			}
			ans[v] = max(ans[v], res);
			if (!st.empty() && st.back() == p) {
				st.pop_back();
				del(p);
			}
		};
		g.dfs(dia[i]);
		dfs(dia[i], -1);
		dfs1(dia[i], -1);
	}
	for (int i = 0; i < n; i++) cout << ans[i] <<"\n";
}

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32