Rapport de résolution pour le problème P4427 [BJOI2018] Somme des profondeurs

1. Énoncé du problème

La tâche principale de ce problème est la suivante : étant donné un arbre ayant le nœud 1 comme racine, il faut répondre à plusieurs requêtes. Chaque requête fournit deux nœuds u, v et un entier k, et nous devons calculer la somme des k-ièmes puissances des profondeurs de tous les nœuds sur le chemin de u à v, le résultat étant pris modulo un certain nombre.

• Profondeur d'un nœud : nombre d'arcs sur le chemin du nœud vers la racine (nœud 1). La profodneur de la racine est 0.
• Échelle des données : le nombre de nœuds n et le nombre de requêtes m peuvent atteindre 300 000, tandis que la valeur maximale de k est 50.

2. Analyse de l'approche

Étape 1 : L'approche la plus intuitive (méthode brute)

Face au problème, l'idée la plus directe est : pour chaque requête (u, v, k), de suivre méthodiquement les étapes suivantes :

1. Trouver le chemin de u à v.
2. Parcourir chaque nœud du chemin.
3. Calculer la k-ième puissance de la profondeur du nœud et l'ajouter à la somme totale.
4. Enfin, afficher la somme totale.

Cette méthode, bien que simple, est inefficace. Lorsque n et m sont grands, trouver le chemin, le parcourir et effectuer les calculs à chaque requête entraînerait un dépassement de temps. La complexité temporelle de cette approche serait d'environ O(N * M), ce qui est inacceptable pour des données de l'ordre de 300 000. Nous avons donc besoin d'un algorithme plus efficace.

Étape 2 : Recherche d'optimisations (prétraitement et "sommes préfixes")

Nous remarquons que la structure de l'arbre est fixe et ne change pas. De plus, la plage de valeurs de k est petite (1 à 50). Cela nous suggère de précalculer et stocker certaines informations utiles afin de pouvoir répondre rapidement aux requêtes.

Cette idée de "prétraitement d'abord, requête rapide ensuite" nous fait penser aux sommes préfixes. Dans un tableau unidimensionnel, pour calculer rapidement la somme d'un intervalle, nous prétraitons la somme de chaque élément jusqu'au point de départ. Sur un arbre, pouvons-nous créer quelque chose de similaire à une "somme préfixe" ?

Oui, bien sûr ! Nous pouvons définir une "somme préfixe sur arbre" :

S(u, k) représente la somme des k-ièmes puissances des profondeurs de tous les nœuds sur le chemin de la racine (nœud 1) au nœud u.

Si nous pouvions précalculer les valeurs de S(u, k) pour tous les nœuds u et toutes les valeurs possibles de k (1 à 50), nous pourrions alors obtenir la réponse à une requête par simple soustraction.

Étape 3 : Dérivation de la formule clé (calcul de la somme du chemin)

Maintenant que nous avons S(u, k), comment l'utiliser pour calculer la somme sur le chemin de u à v ?

Considérons le schéma suivant :

      1 (racine)
      |
      ...
      |
      LCA(u,v)  (plus ancien commun ancestor de u et v)
     /   \
    ...   ...
   /       \
  u         v

• Le chemin de u à v est en fait constitué des deux chemins u -> LCA et v -> LCA concaténés.
• S(u, k) contient la somme des valeurs des nœuds sur le chemin racine -> u.
• S(v, k) contient la somme des valeurs des nœuds sur le chemin racine -> v.

Si nous additionnons directement S(u, k) + S(v, k), que se passe-t-il ?

• Les valeurs des nœuds sur le chemin u -> LCA sont comptées.
• Les valeurs des nœuds sur le chemin v -> LCA sont comptées.
• Les valeurs des nœuds sur le chemin racine -> LCA sont comptées deux fois !

Par conséquent, nous devons soustraire la partie redondante. Naturellement, nous pensons à soustraire S(LCA, k) deux fois. La formule devient : S(u, k) + S(v, k) - 2 * S(LCA, k).

Examinons de plus près quels nœuds cette formule calcule :

• Elle calcule correctement les valeurs des nœuds sur les chemins u -> LCA (sans LCA) et v -> LCA (sans LCA).
• Pour le nœud LCA lui-même : il est compté une fois dans S(u, k), une fois dans S(v, k), et soustrait deux fois dans 2 * S(LCA, k). Finalement, le nœud LCA est compté 1 + 1 - 2 = 0 fois.
• Mais le problème exige que tous les nœuds sur le chemin soient comptés, y compris le nœud LCA !

Ainsi, nous devons ajouter la contribution du nœud LCA lui-même, c'est-à-dire profondeur(LCA)^k. La formule finale est : Réponse = S(u, k) + S(v, k) - 2 * S(LCA, k) + profondeur(LCA)^k.

Cette formule est tout à fait correcte, mais lors de chaque requête, nous devons calculer profondeur(LCA)^k sur place, ce qui est un peu fastidieux. Y a-t-il une manière plus élégante ?

La solution proposée dans l'énoncé est très élégante, dérivons-la :

S(LCA, k) est défini comme la somme sur le chemin racine -> LCA.

S(parent(LCA), k) est défini comme la somme sur le chemin racine -> parent du LCA.

Alors, quelle est la relation entre S(LCA, k) et S(parent(LCA), k) ?

Évidemment : S(LCA, k) = S(parent(LCA), k) + profondeur(LCA)^k.

En réarrangeant, on obtient : profondeur(LCA)^k = S(LCA, k) - S(parent(LCA), k).

Maintenant, substituons cette expression de profondeur(LCA)^k dans la formule finale que nous avons dérivée précédemment :

Réponse = S(u, k) + S(v, k) - 2 * S(LCA, k) + (S(LCA, k) - S(parent(LCA), k))

En simplifiant :

Réponse = S(u, k) + S(v, k) - S(LCA, k) - S(parent(LCA), k)

C'est la formule finale utilisée dans la solution ! Elle est très élégante, tous les termes sont nos valeurs précalculées S, et la requête ne nécessite que des additions et soustractions simples.

3. Détails d'implémentation

Nos idées sont maintenant très claires :

1. Phase de prétraitement : traverser l'arbre entier avec un parcours en profondeur d'abord (DFS).
2. Phase de requête : pour chaque requête, calculer d'abord le LCA, puis appliquer la formule.

Prétraitement (DFS)

Nous devons accomplir toutes les tâches suivantes en un seul parcours DFS :

1. Calculer la profondeur profondeur[u] : c'est simple, profondeur[enfant] = profondeur[parent] + 1.
2. Calculer le parent de chaque nœud parent[u][0] : cela prépare le calcul du LCA.
3. Calculer le tableau de doublement parent[u][j] : utilisé pour calculer rapidement le LCA. parent[u][j] représente le 2^j-ième ancêtre de u, qui peut être dérivé de parent[u][j-1] : parent[u][j] = parent[ parent[u][j-1] ][j-1].
4. Calculer la somme préfixe S(u, k) (dans le code valeur[u][k]) : lorsque nous passons du parent p à l'enfant c, nous pouvons utiliser la valeur S(p, k) déjà calculée pour déterminer S(c, k).

S(c, k) = (S(p, k) + profondeur(c)^k) % mod.

Comme la valeur maximale de k est 50, nous devons calculer 50 valeurs de S pour chaque nœud.

Plus Ancien Commun Ancestor (LCA)

Pour trouver le LCA(u, v), nous avons besoin d'un algorithme efficace. La méthode du doublement est un choix courant et stable, avec une complexité temporelle de O(log N).

Ses étapes de base sont :

1. Faire remonter le nœud le plus profond entre u et v jusqu'à ce qu'il atteigne la même profondeur que l'autre.
2. Si à ce moment les deux nœuds sont identiques, alors c'est le LCA.
3. Sinon, faire remonter simultanément u et v étape par étape jusqu'à ce que leurs parents soient identiques. Ce parent commun est le LCA.

Le processus de "remontée" utilise le tableau parent que nous avons prétraité, permettant de sauter de 2^j étapes à la fois, ce qui est très rapide.

Requête

Pour chaque requête (u, v, k) :

1. Utiliser la méthode du doublement pour calucler L = LCA(u, v).
2. Trouver le parent de L, parent[L][0].
3. Appliquer la formule : réponse = (valeur[u][k] + valeur[v][k] - valeur[L][k] - valeur[parent[L][0]][k]) % mod.
4. Attention : lors du calcul de soustraction avec modulo, le résultat peut devenir négatif. Une pratique standard est (a - b + mod) % mod. L'écriture + 2 * mod dans le code est une façon plus sûre de garantir un résultat positif.

4. Analyse du code

#include<cstdio>
#include<algorithm>
using namespace std;
typedef long long ll;
const ll mod = 998244353;
const int N = 3 * 1e5 + 10;
const int M = 60; // k_max est 50, 60 est suffisant

// --- Définition des variables globales ---
int n, m;
int aretes[2*N], suivants[2*N], debut_liste[N], compteur; // Représentation en liste d'adjacence
ll profondeur[N];                    // profondeur[i] : profondeur du nœud i
int pere[N][22];                     // pere[i][j] : le 2^j-ième ancêtre du nœud i
int marque[N];                       // Marque de visite pour éviter les boucles dans DFS
ll somme[N][M];                     // somme[i][k] : somme des k-ièmes puissances des profondeurs de la racine à i

// --- Ajout d'arête dans la liste d'adjacence ---
inline void ajouter(int u, int v) {
    aretes[++compteur] = v;
    suivants[compteur] = debut_liste[u];
    debut_liste[u] = compteur;
}

// --- Fonction de prétraitement ---
ll puissances[M]; // Tableau temporaire pour calculer rapidement les puissances 1 à 50 d'un nombre (profondeur)
inline void dfs(int u) {
    marque[u] = true;
    // 1. Prétraitement du tableau de doublement pere
    for (int i = 0; pere[u][i]; i++) {
        pere[u][i+1] = pere[pere[u][i]][i];
    }
    // 2. Parcours des enfants
    for (int i = debut_liste[u]; i; i = suivants[i]) {
        if (marque[aretes[i]]) { continue; }
        // 3. Calcul des informations de base de l'enfant : parent et profondeur
        pere[aretes[i]][0] = u;
        profondeur[aretes[i]] = profondeur[u] + 1;
        
        // 4. Calcul de la somme S(v, k) pour l'enfant
        // Calculer d'abord les puissances 1 à 50 de la profondeur de l'enfant
        puissances[0] = 1;
        for (int j = 1; j <= 50; j++) {
            puissances[j] = puissances[j-1] * profondeur[aretes[i]] % mod;
        }
        // En utilisant la somme du parent et la contribution de l'enfant, calculer la somme de l'enfant
        for (int j = 1; j <= 50; j++) {
            somme[aretes[i]][j] = (puissances[j] + somme[u][j]) % mod;
        }
        // Appel récursif pour l'enfant
        dfs(aretes[i]);
    }
}

// --- Calcul de LCA par la méthode du doublement ---
inline int lca(int u, int v) {
    if (profondeur[u] < profondeur[v]) { swap(u, v); } // S'assurer que u est plus profond
    int diff = profondeur[u] - profondeur[v];
    // 1. Faire remonter u pour qu'il atteigne la même profondeur que v
    for (int i = 0; diff; diff >>= 1, i++) {
        if (diff & 1) { u = pere[u][i]; }
    }
    if (u == v) { return u; }
    // 2. Faire remonter u et v simultanément jusqu'à ce que leurs parents soient identiques
    for (int i = 20; i >= 0; i--) {
        if (pere[u][i] != pere[v][i]) {
            u = pere[u][i];
            v = pere[v][i];
        }
    }
    // Retourner le LCA
    return pere[u][0];
}

// --- Fonction principale ---
int main() {
    scanf("%d", &n);
    for (int i = 1, u, v; i < n; i++) {
        scanf("%d%d", &u, &v);
        ajouter(u, v); ajouter(v, u);
    }
    
    dfs(1); // Commencer le prétraitement depuis la racine 1
    
    scanf("%d", &m);
    for (int i = 1, u, v, k; i <= m; i++) {
        scanf("%d%d%d", &u, &v, &k);
        int l = lca(u, v); // Calculer le LCA
        // Appliquer la formule dérivée
        ll reponse = (somme[u][k] + somme[v][k] + 2 * mod - somme[pere[l][0]][k] - somme[l][k]) % mod;
        printf("%lld\n", reponse);
    }
    return 0;
}

Conclusion

Ce problème est un exemple typique qui combine la somme préfixe sur arbre et le Plus Ancien Commun Ancestor (LCA). Les clés de la résolution sont :

1. Utiliser le prétraitement pour optimiser les requêtes, transformant le problème en un problème de "somme préfixe" sur arbre.
2. Dériver correctement la formule de calcul de la somme sur un chemin : Réponse = S(u, k) + S(v, k) - S(LCA, k) - S(parent(LCA), k).
3. Utiliser un parcours DFS efficace pour effectuer tous les prétraitements, y compris la profondeur, les parents, le tableau de doublement et la somme préfixe S.
4. Utiliser des algorithmes efficaces comme la méthode du doublement pour calculer le LCA en O(log N).

Cette combinaison de techniques réduit la complexité de chaque requête de O(N) à O(log N), permettant de résoudre le problème dans les délais impartis.