Submission #1279358

#	Time	Username	Problem	Language	Result	Execution time	Memory
1279358		luvna	LOSTIKS (INOI20_lostiks)	C++20	59 / 100	1854 ms	164624 KiB

Main

#include<bits/stdc++.h>

#define MASK(i) (1 << (i))
#define pub push_back
#define all(v) v.begin(), v.end()
#define compact(v) v.erase(unique(all(v)), end(v))
#define pii pair<int,int>
#define fi first
#define se second
#define endl "\n"
#define sz(v) (int)(v).size()
#define dbg(x) "[" #x " = " << (x) << "]"
#define vi vector<int>

using namespace std;

template<class T> bool minimize(T& a, T b){if(a > b) return a = b, true;return false;}
template<class T> bool maximize(T& a, T b){if(a < b) return a = b, true;return false;}

typedef long long ll;
typedef long double ld;

mt19937_64 rng(chrono::high_resolution_clock::now().time_since_epoch().count());
ll rand(ll l, ll r){return uniform_int_distribution<ll>(l, r)(rng);}

const int MAX = 5e5 + 15;
const int INF = 1e9;

struct DSU{
    int n;
    vector<int> lab;
    stack<pair<int&, int>> history;

    DSU(int n) : n(n), lab(n + 5, -1) {}

    //asce optimization -> peak
    int asc(int u, bool roll){
        return (lab[u] < 0) ? u : (!roll) ? lab[u] = asc(lab[u], roll) : asc(lab[u], roll);
    }

    void join(int u, int v, bool roll){
        u = asc(u, roll), v = asc(v, roll);
        if(u == v) return;
        if(lab[u] > lab[v]) swap(u, v);
        if(roll){
           history.push({lab[u], lab[u]});
           history.push({lab[v], lab[v]});
        }
        lab[u] += lab[v];
        lab[v] = u;
    }

    int snap() {return sz(history);}

    bool same(int u, int v, bool roll){
        return asc(u, roll) == asc(v, roll);
    }

    void rollback(int snapshot){
        assert(snap() >= snapshot);
        while(snap() > snapshot){
            history.top().fi = history.top().se;
            history.pop();
        }
    }
};

int n, m, S, T;
//u: 0, v: 1, key: 2
vector<array<int, 3>> edges;
vector<int> g[MAX];
int tin[MAX], euler[MAX], timer;
int h[MAX];

void dfs(int u, int p){
    euler[tin[u] = ++timer] = u;
    for(int v : g[u]) if(v != p){
        h[v] = h[u] + 1;
        dfs(v, u);
        euler[++timer] = u;
    }
}

int spt[MAX][19];

void prepareLCA(){
    int lg = 31 - __builtin_clz(timer);
    for(int i = 1; i <= timer; i++) spt[i][0] = i;
    for(int j = 1; j <= lg; j++){
        for(int i = 1; i + (1 << j) - 1 <= timer; i++){
            int a = spt[i][j-1];
            int b = spt[i + (1 << (j-1))][j-1];
            spt[i][j] = (h[euler[a]] < h[euler[b]]) ? a : b;
        }
    }
}

int lca(int u, int v){
    int l = tin[u], r = tin[v];
    if(l > r) swap(l, r);
    int k = 31 - __builtin_clz(r - l + 1);
    int a = spt[l][k];
    int b = spt[r - (1 << k) + 1][k];
    return (h[euler[a]] < h[euler[b]]) ? euler[a] : euler[b];
}

int dist(int u, int v){
    return h[u] + h[v] - 2*h[lca(u, v)];
}

//minimum cost to go to vertex U from S, unlock edges in mask
int dp[1 << 20][20][2];

void solve(){
    cin >> n >> S >> T;

    DSU dsu(n);

    for(int i = 1; i < n; i++){
        int u, v, w; cin >> u >> v >> w;
        g[u].push_back(v);
        g[v].push_back(u);
        if(!w){
            dsu.join(u, v, false);
        }
        else{
            edges.push_back({u, v, w});
        }
    }

    dfs(1, -1);
    prepareLCA();

    if(dsu.same(S, T, false)){
        cout << dist(S, T) << endl;
        return;
    }

    m = sz(edges);

    for(int mask = 0; mask < (1 << m); mask++){
        for(int i = 0; i < m; i++){
            for(int j = 0; j < 2; j++){
                dp[mask][i][j] = INF;
            }
        }
    }

    for(int i = 0; i < m; i++){
        if(dsu.same(S, edges[i][2], false)){//get the key without unlock
            for(int j = 0; j < 2; j++){
                if(dsu.same(S, edges[i][j], false)){//S -> key -> U
                    dp[1 << i][i][j] = dist(S, edges[i][2]) + dist(edges[i][2], edges[i][j]);
                }
            }
        }
    }

    //asce optimization 2:
    for(int i = 1; i <= n; i++) dsu.asc(i, false);

    int answer = INF;

    for(int mask = 1; mask < (1 << m); mask++){
        int big_snap = dsu.snap();
        vector<int> lock, unlock;

        for(int i = 0; i < m; i++){
            if(mask >> i & 1){
                unlock.push_back(i);
                dsu.join(edges[i][0], edges[i][1], true);
            }
            else{
                lock.push_back(i);
            }
        }

        //1 case: S, T in same component after unlock edges in mask
        if(dsu.same(S, T, true)){
            for(int i : unlock){
                for(int j = 0; j < 2; j++){
                    minimize(answer, dp[mask][i][j] + dist(edges[i][j], T));
                }
            }
        }

        for(int idF : unlock){
            for(int i = 0; i < 2; i++){
                if(dp[mask][idF][i] == INF) continue;
                for(int idS : lock){
                    for(int j = 0; j < 2; j++){
                        int u, v, key; tie(u, v, key) = {edges[idF][i], edges[idS][j], edges[idS][2]};
                        if(!dsu.same(u, key, true) || !dsu.same(key, v, true)) continue;
                        minimize(dp[mask | (1 << idS)][idS][j], dp[mask][idF][i] + dist(u, key) + dist(key, v));
                    }
                }
            }
        }

        dsu.rollback(big_snap);
    }

        cout << (answer != INF ? answer : -1) << endl;
}

signed main(){
    ios_base::sync_with_stdio(NULL);
    cin.tie(0); cout.tie(0);

    #define task "task"

    if(fopen(task".INP", "r")){
        freopen(task".INP", "r", stdin);
        freopen(task".OUT", "w", stdout);
    }

    int t; t = 1; //cin >> t;
    while(t--) solve();
}

Compilation message (stderr)

Main.cpp: In function 'int main()':
Main.cpp:213:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  213 |         freopen(task".INP", "r", stdin);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
Main.cpp:214:16: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  214 |         freopen(task".OUT", "w", stdout);
      |         ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~

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