[Hackerearth] Tree operations

 

Problem Link    : Tree operations
Category        : Approximate Problem
Contest         : December Circuits'23

#include "bits/stdc++.h"
using namespace std;
#define int     long long int
#define endl    '\n'
 
using pii = pair<int, int>;
 
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
 
int Rand(int l, int r) {
    int rnd = uniform_int_distribution<int>(l, r)(rng);
    return rnd;
}
 
const int maxn = 2e5 + 5;
const int logn = 20;
 
int n;
int k;
int arr[maxn];
vector<vector<int>> graph;
int depth[maxn];
int father[maxn][logn];
 
void dfs(int u, int p, int lvl) {
    depth[u] = lvl;
    for (int i = 1; i < logn; i++) {
        father[u][i] = father[ father[u][i - 1] ][i - 1];
    }
    for (int v : graph[u]) {
        if (v == p) {
            continue;
        }
        father[v][0] = u;
        dfs(v, u, lvl + 1);
    }
}
 
int findLca(int u, int v) {
    if (depth[u] < depth[v]) {
        swap(u, v);
    }
    for (int i = logn - 1; i >= 0; i--) {
        if (depth[ father[u][i] ] >= depth[v]) {
            u = father[u][i];
        }
    }
    if (u == v) {
        return u;
    }
    for (int i = logn - 1; i >= 0; i--) {
        if (father[u][i] != father[v][i]) {
            u = father[u][i];
            v = father[v][i];
        }
    }
    return father[u][0];
}
 
vector<int> up(int u, int lca, bool isForV = false) {
    vector<int> curPath;
    do {
        curPath.emplace_back(u);
        if (u == lca) {
            break;
        }
        u = father[u][0];
    } while (true);
    if (isForV) {
        reverse(curPath.begin(), curPath.end());
        curPath.pop_back();
    }
    return curPath;
}
 
int getEdgesSum(int u, int p) {
    int sum = 0;
    for (int v : graph[u]) {
        if (v == p) {
            continue;
        }
        int add = (arr[u] ^ arr[v]) * (arr[u] | arr[v]);
        sum += add + getEdgesSum(v, u);
    }
    return sum;
}
 
int solve(int &u, int &v, int operations) {
    u = Rand(1, n);
    v = Rand(1, n);
    int lca = findLca(u, v);
    vector<int> pathOfU = up(u, lca);
    vector<int> pathOfV = up(v, lca, true);
    vector<int> path;
    for (int x : pathOfU) {
        path.push_back(x);
    }
    for (int x : pathOfV) {
        path.push_back(x);
    }
    int sz = path.size();
    int temp = arr[ path[0] ];
    for (int i = 0; i < path.size(); i++) {
        int j = (i + 1) % sz;
        arr[ path[i] ] = j == 0 ? temp : arr[ path[j] ];
    }
    int sum = getEdgesSum(1, 0);
    return sum * (k - operations + 1);
}
 
signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr); cout.tie(nullptr);
    cout.precision(12);
 
    bool FILEIO = 1;
    if (FILEIO and fopen("in.txt", "r")) {
        freopen("in.txt", "r", stdin);
    }
 
    cin >> n >> k;
    graph.resize(n + 1);
    for (int i = 1; i <= n; i++) {
        cin >> arr[i];
    }
    for (int e = 1; e < n; e++) {
        int u, v;
        cin >> u >> v;
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    dfs(1, 0, 1);
    int maxi = 0;
    int optimalOperation = 0;
    vector<pii> choosenNodes;
    for (int operations = 1; operations <= k; operations++) {
        int u;
        int v;
        int curSum = solve(u, v, operations);
        choosenNodes.push_back(pii(u, v));
        if (curSum > maxi) {
            optimalOperation = operations;
            maxi = curSum;
        }
    }
    cout << optimalOperation << endl;
    for (int i = 0; i < optimalOperation; i++) {
        cout << choosenNodes[i].first << " " << choosenNodes[i].second << endl;
    }
}