Problème Capoo sur l'arbre

Cette solution combine bisection, décomposition de chaîne d'arbre et arbre de segments persistant pour traiter les requêtes sur des chemins dans un arbre pondéré. L'objectif est de trouver, pour une requête donnée, la distance entre un nœud de départ et le segment continu le plus proche sur le chemin satisfaisant une condition sur les valeurs des nœuds.

Approche algorithmique

La stratégie repose sur la construction d'un arbre de segments persistant indexé par valeur. Plutôt que de stocker les nœuds avec une valeur ≤ val, nous stockons ceux avec une valeur ≥ val en créant des versions de l'arbre de la plus grande valeur à la plus petite. Ainsi, pour une requête de valeur y, la version correspondante contient directement les nœuds satisfaisant la condition.

Pour les opérations de mise à jour (ajout ou soustraction d'unité sur un chemin), nous utilisons la décomposition de chaîne d'arbre (tree chain decomposition) pour convertir le chemin en une collection d'intervalles sur la numérotation DFS (dfn). Ensuite, nous mettons à jour ces intervlales dans l'arbre de segments persistnat en remplaçant les versions correspondantes.

Pour répondre à une requête, nous maintenons des informations sur les segments continus : préfixe maximal, suffixe maximal et longueur maximale d'un segment continu satisfaisant la condition. Ces informations sont fusionnées de manière similaire à la recherche de sous-segment maximum. Par ailleurs, nous utilisons la bisection sur les chemins pour localiser le point de départ du segment, en vérifiant l'existence d'un segment continu de longueur au moins l.

Détails d'implémentation

L'arbre de segments persistant est construit de manière incrémentale. Chaque version correspond à un seuil de valeur. La décomposition de chaîne d'arbre est utilisée pour gérer les mises à jour de chemin et les requêtes de chemin. La fonction de fusion des informations est définie comme suit :

struct SegmentInfo {
    int prefix;
    int suffix;
    int maxLen;
    int length;
};

SegmentInfo merge(const SegmentInfo& left, const SegmentInfo& right) {
    SegmentInfo result;
    result.prefix = (left.prefix == left.length) ? left.prefix + right.prefix : left.prefix;
    result.suffix = (right.suffix == right.length) ? left.suffix + right.suffix : right.suffix;
    result.maxLen = max({left.maxLen, right.maxLen, left.suffix + right.prefix});
    result.length = left.length + right.length;
    return result;
}

Lors du traitement d'une requête de chemin, nous décomposons le chemin en segments via la décomposition de chaîne. Nous parcourons ensuite ces segments pour rechercher un sous-segment continu de longueur au moins l. Si trouvé, la distance est calculée via le plus petit ancêtre commun (LCA). Sinon, la réponse est -1.

Code restructuré

#include <bits/stdc++.h>
using namespace std;

const int MAXN = 100005;
int nodeCount, queryCount;
vector<int> nodeValues;
vector<vector<int>> adjacencyList;

// Variables pour la décomposition de chaîne
int dfsTimer;
vector<int> inTime, outTime, parent, heavyChild, subtreeSize, depth, topNode, nodeId;

void firstDfs(int currentNode, int parentNode) {
    subtreeSize[currentNode] = 1;
    parent[currentNode] = parentNode;
    depth[currentNode] = depth[parentNode] + 1;
    for (int neighbor : adjacencyList[currentNode]) {
        if (neighbor != parentNode) {
            firstDfs(neighbor, currentNode);
            subtreeSize[currentNode] += subtreeSize[neighbor];
            if (!heavyChild[currentNode] || subtreeSize[heavyChild[currentNode]] < subtreeSize[neighbor]) {
                heavyChild[currentNode] = neighbor;
            }
        }
    }
}

void secondDfs(int currentNode, int topValue) {
    topNode[currentNode] = topValue;
    inTime[currentNode] = ++dfsTimer;
    nodeId[dfsTimer] = currentNode;
    if (heavyChild[currentNode]) {
        secondDfs(heavyChild[currentNode], topValue);
    }
    for (int neighbor : adjacencyList[currentNode]) {
        if (neighbor != parent[currentNode] && neighbor != heavyChild[currentNode]) {
            secondDfs(neighbor, neighbor);
        }
    }
}

int computeLca(int u, int v) {
    while (topNode[u] != topNode[v]) {
        if (depth[topNode[u]] < depth[topNode[v]]) swap(u, v);
        u = parent[topNode[u]];
    }
    return depth[u] < depth[v] ? u : v;
}

struct SegmentInfo {
    int prefix, suffix, maxLen, length;
    SegmentInfo() : prefix(0), suffix(0), maxLen(0), length(0) {}
    SegmentInfo(int p, int s, int m, int l) : prefix(p), suffix(s), maxLen(m), length(l) {}
};

SegmentInfo mergeSegment(const SegmentInfo& a, const SegmentInfo& b) {
    SegmentInfo res;
    res.prefix = (a.prefix == a.length) ? a.prefix + b.prefix : a.prefix;
    res.suffix = (b.suffix == b.length) ? a.suffix + b.suffix : b.suffix;
    res.maxLen = max({a.maxLen, b.maxLen, a.suffix + b.prefix});
    res.length = a.length + b.length;
    return res;
}

// Arbre de segments persistant
SegmentInfo segTree[MAXN << 5];
int leftChild[MAXN << 5], rightChild[MAXN << 5];
int roots[MAXN], nodeCounter;

void buildTree(int& node, int l, int r) {
    node = ++nodeCounter;
    segTree[node] = SegmentInfo(0, 0, 0, r - l + 1);
    if (l == r) return;
    int mid = (l + r) >> 1;
    buildTree(leftChild[node], l, mid);
    buildTree(rightChild[node], mid + 1, r);
}

void insertNode(int prevNode, int& newNode, int l, int r, int pos) {
    newNode = ++nodeCounter;
    leftChild[newNode] = leftChild[prevNode];
    rightChild[newNode] = rightChild[prevNode];
    segTree[newNode] = segTree[prevNode];
    if (l == r) {
        segTree[newNode] = SegmentInfo(1, 1, 1, 1);
        return;
    }
    int mid = (l + r) >> 1;
    if (pos <= mid) insertNode(leftChild[prevNode], leftChild[newNode], l, mid, pos);
    else insertNode(rightChild[prevNode], rightChild[newNode], mid + 1, r, pos);
    segTree[newNode] = mergeSegment(segTree[leftChild[newNode]], segTree[rightChild[newNode]]);
}

void copySegment(int& targetRoot, int sourceRoot, int l, int r, int ql, int qr) {
    int temp = targetRoot;
    targetRoot = ++nodeCounter;
    segTree[targetRoot] = segTree[temp];
    leftChild[targetRoot] = leftChild[temp];
    rightChild[targetRoot] = rightChild[temp];
    if (ql <= l && r <= qr) {
        segTree[targetRoot] = segTree[sourceRoot];
        leftChild[targetRoot] = leftChild[sourceRoot];
        rightChild[targetRoot] = rightChild[sourceRoot];
        return;
    }
    int mid = (l + r) >> 1;
    if (ql <= mid) copySegment(leftChild[targetRoot], leftChild[sourceRoot], l, mid, ql, qr);
    if (mid < qr) copySegment(rightChild[targetRoot], rightChild[sourceRoot], mid + 1, r, ql, qr);
    segTree[targetRoot] = mergeSegment(segTree[leftChild[targetRoot]], segTree[rightChild[targetRoot]]);
}

void updatePath(int u, int v, int targetVersion, int sourceVersion) {
    while (topNode[u] != topNode[v]) {
        if (depth[topNode[u]] < depth[topNode[v]]) swap(u, v);
        copySegment(roots[targetVersion], roots[sourceVersion], 1, nodeCount, inTime[topNode[u]], inTime[u]);
        u = parent[topNode[u]];
    }
    if (depth[u] > depth[v]) swap(u, v);
    copySegment(roots[targetVersion], roots[sourceVersion], 1, nodeCount, inTime[u], inTime[v]);
}

SegmentInfo querySegment(int node, int l, int r, int ql, int qr) {
    if (ql <= l && r <= qr) return segTree[node];
    int mid = (l + r) >> 1;
    SegmentInfo ans;
    if (ql <= mid) ans = mergeSegment(ans, querySegment(leftChild[node], l, mid, ql, qr));
    if (mid < qr) ans = mergeSegment(ans, querySegment(rightChild[node], mid + 1, r, ql, qr));
    return ans;
}

void processQuery(int u, int v, int version, int minLen) {
    int startNode = u;
    vector<pair<int, int>> segments, reverseSegments;
    while (topNode[u] != topNode[v]) {
        if (depth[topNode[u]] > depth[topNode[v]]) {
            segments.emplace_back(inTime[u], inTime[topNode[u]]);
            u = parent[topNode[u]];
        } else {
            reverseSegments.emplace_back(inTime[topNode[v]], inTime[v]);
            v = parent[topNode[v]];
        }
    }
    segments.emplace_back(inTime[u], inTime[v]);
    for (int i = reverseSegments.size() - 1; i >= 0; --i) {
        segments.emplace_back(reverseSegments[i]);
    }

    SegmentInfo currentInfo;
    int targetPosition = 0;
    vector<pair<int, int>> processedSegments;
    for (auto seg : segments) {
        int l = seg.first, r = seg.second;
        SegmentInfo segInfo;
        if (l < r) segInfo = querySegment(roots[version], 1, nodeCount, l, r);
        else segInfo = querySegment(roots[version], 1, nodeCount, r, l);
        if (l > r) swap(segInfo.prefix, segInfo.suffix);

        if (currentInfo.suffix + segInfo.prefix >= minLen && currentInfo.suffix) {
            int remaining = currentInfo.suffix;
            for (int j = processedSegments.size() - 1; j >= 0; --j) {
                int segLen = abs(processedSegments[j].first - processedSegments[j].second) + 1;
                if (remaining > segLen) remaining -= segLen;
                else {
                    auto [start, end] = processedSegments[j];
                    if (start < end) targetPosition = nodeId[end - remaining + 1];
                    else targetPosition = nodeId[end + remaining - 1];
                    break;
                }
            }
            break;
        }

        currentInfo = mergeSegment(currentInfo, segInfo);
        processedSegments.emplace_back(seg);

        if (segInfo.maxLen >= minLen) {
            if (l < r) {
                int low = l, high = r;
                while (low <= high) {
                    int mid = (low + high) >> 1;
                    if (querySegment(roots[version], 1, nodeCount, l, mid).maxLen >= minLen) high = mid - 1;
                    else low = mid + 1;
                }
                targetPosition = nodeId[high + 1 - minLen + 1];
            } else {
                int low = r, high = l;
                while (low <= high) {
                    int mid = (low + high) >> 1;
                    if (querySegment(roots[version], 1, nodeCount, mid, l).maxLen >= minLen) low = mid + 1;
                    else high = mid - 1;
                }
                targetPosition = nodeId[low - 1 + minLen - 1];
            }
            break;
        }
    }

    if (targetPosition) {
        int lcaNode = computeLca(startNode, targetPosition);
        int distance = depth[startNode] + depth[targetPosition] - 2 * depth[lcaNode];
        cout << distance << '\n';
    } else {
        cout << -1 << '\n';
    }
}

void solve() {
    cin >> nodeCount >> queryCount;
    nodeValues.resize(nodeCount + 1);
    adjacencyList.resize(nodeCount + 1);
    inTime.resize(nodeCount + 1);
    outTime.resize(nodeCount + 1);
    parent.resize(nodeCount + 1);
    heavyChild.resize(nodeCount + 1);
    subtreeSize.resize(nodeCount + 1);
    depth.resize(nodeCount + 1);
    topNode.resize(nodeCount + 1);
    nodeId.resize(nodeCount + 1);

    vector<vector<int>> valueNodes(nodeCount + 1);
    for (int i = 1; i <= nodeCount; ++i) {
        cin >> nodeValues[i];
        valueNodes[nodeValues[i]].push_back(i);
    }
    for (int i = 1; i < nodeCount; ++i) {
        int u, v; cin >> u >> v;
        adjacencyList[u].push_back(v);
        adjacencyList[v].push_back(u);
    }

    dfsTimer = 0;
    firstDfs(1, 0);
    secondDfs(1, 1);

    buildTree(roots[nodeCount + 1], 1, nodeCount);
    for (int val = nodeCount; val >= 1; --val) {
        roots[val] = roots[val + 1];
        for (int pos : valueNodes[val]) {
            insertNode(roots[val], roots[val], 1, nodeCount, inTime[pos]);
        }
    }

    for (int i = 0; i < queryCount; ++i) {
        int op; cin >> op;
        if (op == 1) {
            int u, v, x; cin >> u >> v >> x;
            updatePath(u, v, x, x + 1);
        } else if (op == 2) {
            int u, v, x; cin >> u >> v >> x;
            updatePath(u, v, x + 1, x);
        } else {
            int u, v, l, y; cin >> u >> v >> l >> y;
            processQuery(u, v, y, l);
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    solve();
    return 0;
}