Submission #624728

#	Time	Username	Problem	Language	Result	Execution time	Memory
624728		Tigerpants	Ideal city (IOI12_city)	C++17	100 / 100	219 ms	29360 KiB

city

#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <algorithm>
#include <numeric>
#include <iomanip>
#include <string>
#include <deque>
#include <random>

using namespace std;

typedef long long int ll;
typedef vector<ll> vi;
typedef vector<vi> vvi;
typedef multiset<ll> multisi;
typedef set<ll> si;
typedef vector<si> vsi;
typedef pair<ll, ll> pi;
typedef vector<pi> vpi;
typedef set<pi> spi;
typedef map<pi, ll> mpii;

ll N;
vpi cells;

vi x_to_auxillary_node;
vi size_of_x_auxillary_node;
vi y_to_auxillary_node;
vi size_of_y_auxillary_node;

si x_cells;
si y_cells;
vsi x_cells_graph;
vsi y_cells_graph;

// distance in three between nodes i and j is: dist[i] + dist[j] - 2 * dist[LCA(i, j)]
// compute dist in O(n)
// compute LCA in O(nlogn)
// sum = sum_dist(w[i] * dist[i] * M) - 2 * sum_LCA(dist[LCA]).
// so we sum over the distance to the LCA of every pair of nodes
// if node has ≥2 children: 1 + sum(subgraphsize[i] * (1 + sum_subgraphsize - subgraphsize[i])) // node <-> node, subgraph <-> node, subgraph <-> subgraph
// else node has < 2 children: LCA = subgraphsize(child) + 1
// count the number of nodes a node is LCA for: number of pairs of nodes in subgraph: (#nodes in subgraph)^2. (sum(nodes_in_subgraph for children) + nodes_in_self)

bool cell_exists(ll x, ll y) {
    vpi::iterator cell = lower_bound(cells.begin(), cells.end(), make_pair(x, y));
    if (cell == cells.end()) {
        return false;
    }
    return (*cell == make_pair(x, y));
}

ll find_par_x(ll x) {
    ll y = x;
    while (x != x_to_auxillary_node[x]) {
        x = x_to_auxillary_node[x];
    }
    x_to_auxillary_node[y] = x;
    return x;
}

ll find_par_y(ll x) {
    ll y = x;
    while (x != y_to_auxillary_node[x]) {
        x = y_to_auxillary_node[x];
    }
    y_to_auxillary_node[y] = x;
    return x;
}

void x_combine(ll a, ll b) {
    a = find_par_x(a);
    b = find_par_x(b);
    if (a == b) {
        return;
    }
    if (size_of_x_auxillary_node[a] > size_of_x_auxillary_node[b]) {
        int temp = a;
        a = b;
        b = temp;
    }
    x_to_auxillary_node[a] = b;
    size_of_x_auxillary_node[b] += size_of_x_auxillary_node[a];
    x_cells.erase(a);
}

void y_combine(ll a, ll b) {
    a = find_par_y(a);
    b = find_par_y(b);
    if (a == b) {
        return;
    }
    if (size_of_y_auxillary_node[a] > size_of_y_auxillary_node[b]) {
        int temp = a;
        a = b;
        b = temp;
    }
    y_to_auxillary_node[a] = b;
    size_of_y_auxillary_node[b] += size_of_y_auxillary_node[a];
    y_cells.erase(a);
}

vi subgraph_size; // number cells in subtree
vi dist_from_root;

ll sum_pairwise_distances = 0;

// sum_{i}^N{sum_{j}^N{dist[i][j]}} = N * sum_{i}^N(dist[i]) - sum_{i}^N{sum_{pair nodes with LCA i}{dist[i]}}
// sum_{i}^N{sum_{pair nodes with LCA i}{dist[i]}} = sum_{i}^N{dist[i] * (pair nodes with LCA i)}
// pair nodes with LCA i = (pairs of nodes in different child subtrees of i) + (pairs of cell in node i and in strict subtree of i) + (pairs of cells in node i)
// pair nodes with LCA i = sum_{child j of i}{(size[subtree(j)] * (sum_{child k != j of i}{size[subtree(j)]}))} / 2 + sum_{child j of i}{size[subtree(j)] * size[i]} + size[i]^2
// LCA = sum(subgraph_size[*child] * (subgraph_size[cur] - subgraph_size[*child])) / 2  + subgraph_size[cur] - size_of_x_auxillary_node[cur] + size_of_x_auxilary_node[cur]**2;

void x_DFS(ll cur, ll par) {
    //cout << cur << ": (";
    subgraph_size[cur] = size_of_x_auxillary_node[cur];
    for (si::iterator child = x_cells_graph[cur].begin(); child != x_cells_graph[cur].end(); child++) {
        if (*child != par) {
            dist_from_root[*child] = dist_from_root[cur] + 1;
            x_DFS(*child, cur);
            subgraph_size[cur] += subgraph_size[*child];
        }
    }
    
    // calculate pair of nodes with LCA i
    ll twoLCA = subgraph_size[cur] * subgraph_size[cur];
    for (si::iterator child = x_cells_graph[cur].begin(); child != x_cells_graph[cur].end(); child++) {
        if (*child != par) {
            twoLCA -= subgraph_size[*child] * subgraph_size[*child];
        }
    }
    twoLCA *= dist_from_root[cur];

    sum_pairwise_distances += dist_from_root[cur] * N * size_of_x_auxillary_node[cur];
    sum_pairwise_distances -= twoLCA;

    //cout << ")->" << sum_pairwise_distances << " | ";
}

void y_DFS(ll cur, ll par) {
    //cout << cur << ": (";
    subgraph_size[cur] = size_of_y_auxillary_node[cur];
    for (si::iterator child = y_cells_graph[cur].begin(); child != y_cells_graph[cur].end(); child++) {
        if (*child != par) {
            dist_from_root[*child] = dist_from_root[cur] + 1;
            y_DFS(*child, cur);
            subgraph_size[cur] += subgraph_size[*child];
        }
    }
    
    // calculate pair of nodes with LCA i
    ll twoLCA = subgraph_size[cur] * subgraph_size[cur];
    for (si::iterator child = y_cells_graph[cur].begin(); child != y_cells_graph[cur].end(); child++) {
        if (*child != par) {
            twoLCA -= subgraph_size[*child] * subgraph_size[*child];
        }
    }
    twoLCA *= dist_from_root[cur];

    sum_pairwise_distances += dist_from_root[cur] * N * size_of_y_auxillary_node[cur];
    sum_pairwise_distances -= twoLCA;

    //cout << ")->" << sum_pairwise_distances << " | ";
}

int DistanceSum(int N_, int* X_, int* Y_) {
    N = N_;

    cells.resize(N);

    for (int i = 0; i < N; i++) {
        cells[i] = make_pair(X_[i], Y_[i]);
    }

    sort(cells.begin(), cells.end());

    x_to_auxillary_node.resize(N);
    size_of_x_auxillary_node.resize(N);
    y_to_auxillary_node.resize(N);
    size_of_y_auxillary_node.resize(N);
    x_cells_graph.resize(N);
    y_cells_graph.resize(N);
    for (int i = 0; i < N; i++) {
        x_to_auxillary_node[i] = i;
        size_of_x_auxillary_node[i] = 1;
        y_to_auxillary_node[i] = i;
        size_of_y_auxillary_node[i] = 1;
        x_cells.insert(i);
        y_cells.insert(i);
    }
    
    // create nodes in auxillary graphs
    for (int i = 0; i < N; i++) {
        if (cell_exists(cells[i].first + 1, cells[i].second)) {
            x_combine(i, lower_bound(cells.begin(), cells.end(), make_pair(cells[i].first + 1, cells[i].second)) - cells.begin());
        }
        if (cell_exists(cells[i].first - 1, cells[i].second)) {
            x_combine(i, lower_bound(cells.begin(), cells.end(), make_pair(cells[i].first - 1, cells[i].second)) - cells.begin());
        }
        if (cell_exists(cells[i].first, cells[i].second + 1)) {
            y_combine(i, lower_bound(cells.begin(), cells.end(), make_pair(cells[i].first, cells[i].second + 1)) - cells.begin());
        }
        if (cell_exists(cells[i].first, cells[i].second - 1)) {
            y_combine(i, lower_bound(cells.begin(), cells.end(), make_pair(cells[i].first, cells[i].second - 1)) - cells.begin());
        }
    }

    // create edges in auxillary graphs
    for (int i = 0; i < N; i++) {
        if (cell_exists(cells[i].first + 1, cells[i].second)) {
            y_cells_graph[find_par_y(i)].insert(find_par_y(lower_bound(cells.begin(), cells.end(), make_pair(cells[i].first + 1, cells[i].second)) - cells.begin()));
        }
        if (cell_exists(cells[i].first - 1, cells[i].second)) {
            y_cells_graph[find_par_y(i)].insert(find_par_y(lower_bound(cells.begin(), cells.end(), make_pair(cells[i].first - 1, cells[i].second)) - cells.begin()));
        }
        if (cell_exists(cells[i].first, cells[i].second + 1)) {
            x_cells_graph[find_par_x(i)].insert(find_par_x(lower_bound(cells.begin(), cells.end(), make_pair(cells[i].first, cells[i].second + 1)) - cells.begin()));
        }
        if (cell_exists(cells[i].first, cells[i].second - 1)) {
            x_cells_graph[find_par_x(i)].insert(find_par_x(lower_bound(cells.begin(), cells.end(), make_pair(cells[i].first, cells[i].second - 1)) - cells.begin()));
        }
    }

    // now we have the nodes and edges in our axillary graphs
    // nodes are in the si x_cells
    // edges are in the vsi x_cells_graph
    // the auxillary graphs are trees
    // we root the trees and find the sum of the nodes in the sumgraph, as well as the distance to each node
    // this we do through dfs

    /*
    for (auto itr : x_to_auxillary_node) {
        cout << itr << " ";
    }
    cout << endl;
    for (auto itr : size_of_x_auxillary_node) {
        cout << itr << " ";
    }
    cout << endl;
    for (auto itr : y_to_auxillary_node) {
        cout << itr << " ";
    }
    cout << endl;
    for (auto itr : size_of_y_auxillary_node) {
        cout << itr << " ";
    }
    cout << endl;
    for (auto itr : y_cells) {
        cout << itr << " ";
    }
    cout << endl;
    for (auto itr : x_cells) {
        cout << itr << " ";
    }
    cout << endl;
    for (int i = 0; i < N; i++) {
        if (y_cells_graph[i].size()) {
            cout << i << ": ";
            for (auto itr : y_cells_graph[i]) {
                cout << itr << " ";
            }
            cout << endl;
        }
    }
    for (int i = 0; i < N; i++) {
        if (x_cells_graph[i].size()) {
            cout << i << ": ";
            for (auto itr : x_cells_graph[i]) {
                cout << itr << " ";
            }
            cout << endl;
        }
    }
    */

    subgraph_size.resize(N);
    dist_from_root.resize(N);

    dist_from_root[*x_cells.begin()] = 0;
    x_DFS(*x_cells.begin(), -1);

    //cout << sum_pairwise_distances << endl;

    dist_from_root[*y_cells.begin()] = 0;
    y_DFS(*y_cells.begin(), -1);

    ll big_mod = 1000000000;
    //cout << sum_pairwise_distances % big_mod << endl;
    return sum_pairwise_distances % big_mod;

    return 0;
}

• Subtask 1: 11 / 11
• Subtask 2: 21 / 21
• Subtask 3: 23 / 23
• Subtask 4: 45 / 45