Submission #503006

#	Time	Username	Problem	Language	Result	Execution time	Memory
503006		Victor	Fun Tour (APIO20_fun)	C++17	100 / 100	288 ms	21212 KiB

fun

// #pragma GCC target ("avx,avx2,fma")
// #pragma GCC optimize ("Ofast,inline") // O1 - O2 - O3 - Os - Ofast
// #pragma GCC optimize ("unroll-loops")
#include "fun.h"

#include <bits/stdc++.h>

using namespace std;

#define rep(i, a, b) for (int i = (a); i < (b); ++i)
#define per(i, a, b) for (int i = (b - 1); i >= (a); --i)
#define trav(a, x) for (auto& a : x)

#define all(x) x.begin(), x.end()
#define sz(x) x.size()
#define pb push_back
#define debug(x) cout << #x << " = " << x << endl

#define umap unordered_map
#define uset unordered_set

typedef pair<int, int> ii;
typedef pair<int, ii> iii;
typedef vector<int> vi;
typedef vector<ii> vii;
typedef vector<vi> vvi;

typedef long long ll;
typedef pair<ll, ll> pll;
typedef vector<ll> vll;
typedef vector<pll> vpll;

const int INF = 1'000'000'007;

int n, q, sub_size[100001], dist[100001];
int abc = 0;
vii subtrees[3];

std::vector<int> createFunTour(int N, int Q) {
    if (N == 2) {
        vi vec = {0, 1};
        return vec;
    }

    n = N;
    q = Q;
    int centroid = 0, mn = INF;
    rep(i, 1, n) {
        sub_size[i] = attractionsBehind(0, i);
        if (sub_size[i] >= (n + 1) / 2 && sub_size[i] < mn) {
            centroid = i;
            mn = sub_size[i];
        }
    }

    int cnt = 0;
    vi neighbours;
    rep(i, 0, n) if (i != centroid) {
        dist[i] = hoursRequired(centroid, i);
        if (dist[i] == 1) {
            neighbours.pb(i);
            ++cnt;
        }
    }

    rep(i, 0, n) if (i != centroid) {
        per(j, 0, cnt) {
            bool yes = 1;
            if (j) yes = dist[i] == hoursRequired(neighbours[j], i) + 1;
            if (yes) {
                subtrees[j].emplace_back(dist[i], i);
                break;
            }
        }
    }

    int add_last=-1;
    vi ans;
    rep(i, 0, cnt) {
        sort(all(subtrees[i]));
        if(sz(subtrees[i])==(n+1)/2) {
            subtrees[i].erase(subtrees[i].begin());
            add_last=neighbours[i];
        }
    }

    int prev = -1, prev_val = INF;
    if (cnt == 3) {
        while (1) {
            vii sizes;
            rep(i, 0, cnt) sizes.emplace_back(sz(subtrees[i]), i);
            sort(all(sizes));

            if (sizes[0].first + sizes[1].first == sizes[2].first) {
                int node, sub, val = -1;

                int smallest = sizes[0].second, small = sizes[1].second;

                rep(i, 0, cnt) if (prev != i && subtrees[i].back().first > val) {
                    tie(val, node) = subtrees[i].back();
                    sub = i;
                }

                if (((prev == small && sub == smallest) || (prev == smallest && sub == small)) && subtrees[sub].back().first > prev_val) {
                    subtrees[sub].pop_back();
                    ans.pb(node);
                    prev = sub;
                }

                if (smallest < small) swap(small, smallest);
                trav(val, subtrees[smallest]) subtrees[small].pb(val);

                if (!small) {
                    if (smallest == 1) {
                        subtrees[0].swap(subtrees[2]);
                        subtrees[1].swap(subtrees[2]);

                    } else
                        subtrees[0].swap(subtrees[1]);
                }

                if (prev == small || prev == smallest)
                    prev = 1;
                else
                    prev = 0;
                break;
            }

            int node, sub, val = -1;
            rep(i, 0, cnt) if (prev != i && subtrees[i].back().first > val) {
                tie(val, node) = subtrees[i].back();
                sub = i;
            }

            subtrees[sub].pop_back();
            ans.pb(node);
            prev = sub;
            prev_val = val;
        }
    }

    rep(i, 0, 2) sort(all(subtrees[i]));
    if (prev == -1) prev = sz(subtrees[0]) > sz(subtrees[1]);
    while (1) {
        prev ^= 1;
        if (subtrees[prev].empty()) break;
        ans.pb(subtrees[prev].back().second);
        subtrees[prev].pop_back();
    }

    ans.pb(centroid);
    if(add_last!=-1)ans.pb(add_last);
    // trav(val,ans)cout<<"u = "<<val<<endl;

    return ans;
}

/*
static void wrongAnswer(std::string message) {
    printf("WA: %s\n", message.c_str());
    exit(0);
}

namespace tree_helper {

static int N;
static int logN;
static std::vector<std::vector<int>> parent;
static std::vector<int> depth;
static std::vector<int> subtree_size;

static void dfs(
    const std::vector<std::vector<int>>& adj_list,
    int current_node, int parent_node) {
    parent[0][current_node] = parent_node;
    subtree_size[current_node] = 1;
    for (int i = 0; i < static_cast<int>(adj_list[current_node].size()); ++i) {
        const int next_node = adj_list[current_node][i];
        if (next_node != parent_node) {
            depth[next_node] = depth[current_node] + 1;
            dfs(adj_list, next_node, current_node);
            subtree_size[current_node] += subtree_size[next_node];
        }
    }
}

static void initializeTree(const std::vector<std::vector<int>>& adj_list) {
    N = static_cast<int>(adj_list.size());

    depth = std::vector<int>(N, 0);
    subtree_size = std::vector<int>(N, 0);
    for (logN = 0; (1 << logN) < N; ++logN) {
    }
    parent = std::vector<std::vector<int>>(logN, std::vector<int>(N, 0));

    dfs(adj_list, 0, 0);
    for (int i = 1; i < logN; ++i) {
        for (int j = 0; j < N; ++j) {
            parent[i][j] = parent[i - 1][parent[i - 1][j]];
        }
    }
}

static int getLowestCommonAncestor(int X, int Y) {
    if (depth[X] < depth[Y]) {
        std::swap(X, Y);
    }
    for (int i = logN - 1; i >= 0; --i) {
        if (depth[parent[i][X]] >= depth[Y]) {
            X = parent[i][X];
        }
    }
    if (X == Y) {
        return X;
    }
    for (int i = logN - 1; i >= 0; --i) {
        if (parent[i][X] != parent[i][Y]) {
            X = parent[i][X];
            Y = parent[i][Y];
        }
    }
    return parent[0][X];
}

static int getDistance(int X, int Y) {
    return depth[X] + depth[Y] - 2 * depth[getLowestCommonAncestor(X, Y)];
}

static int attractionsBehind(int X, int Y) {
    if (X == Y) {
        return N;
    }
    for (int i = logN - 1; i >= 0; --i) {
        if (depth[parent[i][X]] > depth[Y]) {
            X = parent[i][X];
        }
    }
    if (Y == parent[0][X]) {
        return N - subtree_size[X];
    }
    return subtree_size[Y];
}

static void checkFunTour(const std::vector<int>& fun_tour) {
    if (static_cast<int>(fun_tour.size()) != N) {
        wrongAnswer("Invalid size");
    }

    std::vector<bool> visited_attractions(N, false);
    for (int i = 0; i < N; ++i) {
        if (fun_tour[i] < 0 || fun_tour[i] >= N) {
            wrongAnswer("Invalid index");
        }
        if (visited_attractions[fun_tour[i]]) {
            wrongAnswer("Repeated index");
        }
        visited_attractions[fun_tour[i]] = true;
    }

    int last_travel_time = getDistance(fun_tour[0], fun_tour[1]);
    for (int i = 2; i < N; ++i) {
        int travel_time = getDistance(fun_tour[i - 1], fun_tour[i]);
        if (travel_time > last_travel_time) {
            cout << "i = " << i << endl;
            wrongAnswer("Tour is not fun");
        }
        last_travel_time = travel_time;
    }
}

}  // namespace tree_helper

static int N, Q;

int hoursRequired(int X, int Y) {
    if (--Q < 0) {
        wrongAnswer("Too many queries");
    }
    if (X < 0 || X >= N || Y < 0 || Y >= N) {
        wrongAnswer("Invalid index");
    }
    return tree_helper::getDistance(X, Y);
}

int attractionsBehind(int X, int Y) {
    if (--Q < 0) {
        wrongAnswer("Too many queries");
    }
    if (X < 0 || X >= N || Y < 0 || Y >= N) {
        wrongAnswer("Invalid index");
    }
    return tree_helper::attractionsBehind(X, Y);
}

int main() {
    assert(2 == scanf("%d %d", &N, &Q));

    std::vector<std::vector<int>> adj_list(N);
    for (int i = 0; i < N - 1; ++i) {
        int A, B;
        assert(2 == scanf("%d %d", &A, &B));
        adj_list[A].push_back(B);
        adj_list[B].push_back(A);
    }
    tree_helper::initializeTree(adj_list);

    std::vector<int> fun_tour = createFunTour(N, Q);
    tree_helper::checkFunTour(fun_tour);

    for (int i = 0; i < N; ++i) {
        printf("%d%c", fun_tour[i], " \n"[i == N - 1]);
    }
    return 0;
}*/

Compilation message (stderr)

fun.cpp: In function 'std::vector<int> createFunTour(int, int)':
fun.cpp:81:27: warning: comparison of integer expressions of different signedness: 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   81 |         if(sz(subtrees[i])==(n+1)/2) {
      |                           ^
fun.cpp:122:17: warning: 'sub' may be used uninitialized in this function [-Wmaybe-uninitialized]
  122 |                 if (prev == small || prev == smallest)
      |                 ^~

• Subtask 1: 10 / 10
• Subtask 2: 16 / 16
• Subtask 3: 21 / 21
• Subtask 4: 19 / 19
• Subtask 5: 34 / 34