Submission #999272

#	Time	Username	Problem	Language	Result	Execution time	Memory
999272		vjudge1	LOSTIKS (INOI20_lostiks)	C++17	100 / 100	1225 ms	492096 KiB

Main

#include <bits/stdc++.h>

using namespace std;

#define int long long
#define inf 0x3F3F3F3F3F3F3F3F

struct DSU
{
	vector<int> e;
	void init(int n)
	{
		e.assign(n + 1, -1);
	}
	int get(int x)
	{
		if (e[x] < 0) return x;
		return e[x] = get(e[x]);
	}
	void unite(int x, int y)
	{
		x = get(x), y = get(y);
		if (x == y) return;
		if (e[x] > e[y]) swap(x, y);
		e[x] += e[y];
		e[y] = x;
	}
};

const int MXN = 1e6 + 5;
const int LOG = 20;

int n, S, T;
vector<array<int, 2>> adj[MXN];
int p[LOG][MXN], dep[MXN], in[MXN], out[MXN], tim;
int b[MXN], need[MXN];
int dp[20][(1LL << 20)];
vector<array<int, 3>> m;

void dfs(int a)
{
	in[a] = ++tim;
	for (array<int, 2> &v : adj[a])
	{
		if (v[0] == p[0][a]) continue;
		p[0][v[0]] = a;
		dep[v[0]] = dep[a] + 1;
		need[v[0]] = need[a];
		if (v[1])
		{
			int ind = lower_bound(m.begin(), m.end(), array<int, 3> {a, v[0], v[1]}) - m.begin();
			need[v[0]] = need[a] | (1LL << ind);
		}
		dfs(v[0]);
	}
	out[a] = tim;
}

int anc(int u, int v)
{
	return in[u] <= in[v] && out[v] <= out[u];
}
int LCA(int u, int v)
{
	if (anc(u, v)) return u;
	if (anc(v, u)) return v;
	for (int i = LOG - 1; i >= 0; i--)
	{
		if (anc(p[i][u], v)) continue;
		u = p[i][u];
	}
	return p[0][u];
}
int DIST(int u, int v)
{
	return dep[u] + dep[v] - 2 * dep[LCA(u, v)];
}

signed main()
{
	ios::sync_with_stdio(0);
	cin.tie(0);
	cin >> n >> S >> T;
	DSU dsu;
	dsu.init(n);
	for (int i = 1; i < n; i++)
	{
		int u, v, w;
		cin >> u >> v >> w;
		adj[u].push_back({v, w});
		adj[v].push_back({u, w});
		if (w) m.push_back({u, v, w});
		if (!w) dsu.unite(u, v);
	}
	sort(m.begin(), m.end());
	p[0][S] = S;
	dfs(S);
	for (array<int, 3> &i : m) if (p[0][i[0]] == i[1]) swap(i[0], i[1]);
	dfs(S);
	for (int i = 1; i < LOG; i++) for (int j = 1; j <= n; j++) p[i][j] = p[i - 1][p[i - 1][j]];
	for (int i = 0; i < m.size(); i++) for (int j = 0; j < (1LL << (int)m.size()); j++) dp[i][j] = inf;
	for (int i = 0; i < m.size(); i++)
	{
		b[i] = (need[m[i][0]] ^ need[m[i][2]] ^ need[LCA(m[i][0], m[i][2])]);
	}
	if (dsu.get(S) == dsu.get(T)) 
	{
		cout << DIST(S, T) << '\n';
		return 0;
	}
	for (int msk = 1; msk < (1LL << (int)(m.size())); msk++)
	{
		if (__builtin_popcount(msk) == 1)
		for (int i = 0; i < m.size(); i++)
		{
			if (!b[i] && (msk >> i & 1)) 
			{
				dp[i][msk] = DIST(S, m[i][2]) + DIST(m[i][2], m[i][0]);
			}
		}
		for (int i = 0; i < m.size(); i++)
		{
			if (dp[i][msk] >= inf) continue;
			for (int j = 0; j < m.size(); j++)
			{
				if ((msk >> j & 1)) continue;
				if (((msk ^ b[j]) & b[j])) continue;
				dp[j][msk | (1LL << j)] = min(dp[j][msk | (1LL << j)], dp[i][msk] + DIST(m[i][0], m[j][2]) + DIST(m[j][2], m[j][0]));
			}
		}
	}
	int res = inf;
	for (int msk = 1; msk < (1LL << (int)m.size()); msk++)
	{
		for (int i = 0; i < m.size(); i++)
		{
			if (dp[i][msk] >= inf) continue;
			if (((msk ^ need[T]) & need[T])) continue;
			res = min(res, dp[i][msk] + DIST(m[i][0], T));
 		}
	}
	cout << (res >= inf ? -1 : res) << '\n';
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:101:20: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::array<long long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  101 |  for (int i = 0; i < m.size(); i++) for (int j = 0; j < (1LL << (int)m.size()); j++) dp[i][j] = inf;
      |                  ~~^~~~~~~~~~
Main.cpp:102:20: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::array<long long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  102 |  for (int i = 0; i < m.size(); i++)
      |                  ~~^~~~~~~~~~
Main.cpp:114:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::array<long long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  114 |   for (int i = 0; i < m.size(); i++)
      |                   ~~^~~~~~~~~~
Main.cpp:121:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::array<long long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  121 |   for (int i = 0; i < m.size(); i++)
      |                   ~~^~~~~~~~~~
Main.cpp:124:22: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::array<long long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  124 |    for (int j = 0; j < m.size(); j++)
      |                    ~~^~~~~~~~~~
Main.cpp:135:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::array<long long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  135 |   for (int i = 0; i < m.size(); i++)
      |                   ~~^~~~~~~~~~

• Subtask 1: 23 / 23
• Subtask 2: 36 / 36
• Subtask 3: 41 / 41