Robię to zadanie: http://www.usaco.org/index.php?page=viewproblem2&cpid=968
Używam przygotowuje dane pod LCA, sprawdzam czy na ścieżce 2 wierzchołków do LCA jest szukany rodzaj. Niestety mój program przechodzi tylko jeden test, więc albo nie rozumiem treści zadania, albo mam jakieś błedy w kodzie. Wiadomość czy dany wierzochłek ma rodzaj "H" czy "G" trzymam w masce bitowej.
Mój program zwraca że na ścieżce wystąpił dany rodzaj pomimo tego że go nie ma. (przynajmniej w jednym z testów tak jest, niestety testy mają n = 1000 więc nie mam jak ich analizować)

Kopiuj
#include <bits/stdc++.h>
    
using namespace std;
 
#define ll long long
#define ull unsigned long long
 
const ll INF = 1e9 + 7, MAXN = 1e5 + 7, MAXLOG = 17;

vector<vector<int>> graph(MAXN), jumps(MAXN, vector<int>(MAXLOG)), milk(MAXN, vector<int>(MAXLOG));
vector<int> depht(MAXN);

void setIO(string s){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    if (s != "-"){
        freopen((s + ".in").c_str(), "r", stdin);
        freopen((s + ".out").c_str(), "w", stdout);
    }
}

void dfs(int p, int v){
    jumps[v][0] = p;
    depht[v] = depht[p] + 1;
    for (auto child : graph[v])
        if (child != p)
            dfs(v, child);
}

bool lca(int x, int y, int f){
    int found = 0;

    //zawsze x jest nizej
    if (depht[x] < depht[y])
        swap(x, y);
    
    //wyrownujemy poziomy
    for (int i = MAXLOG-1; i >= 0; i--){
        if (depht[x] - (1 << i) >= depht[y]){
            found |= milk[x][i];
            x = jumps[x][i];
        }
    }

    //wyrownujac poziomy trafilismy do tego samego wierzochlka
    if (x == y)
        // spotkalismy 2 lub szukany  lub dla bezpieczenstwa sprawdzam 
        return found == 3 || found == f || milk[x][0] == f;
    
    //idziemy synow LCA
    for (int i = MAXLOG-1; i >= 0; i--){
        if (jumps[x][i] != jumps[y][i]){
            found |= milk[x][i];
            x = jumps[x][i];
            found |= milk[y][i];
            y = jumps[y][i];
        }
    }

    return found == 3 || found == f || milk[x][1] == f; // ojciec moze miec szukany rodzaj
}

int main(){
    setIO("1");
    int n, m;
    cin >> n >> m;

    char breed;
    for (int i = 1; i <= n; i++){
        cin >> breed;
        milk[i][0] = (breed == 'H') ? 1 : 2;
    }

    int a, b;
    for (int i = 0; i < n-1; i++){
        cin >> a >> b;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }

    dfs(1, 1);

    for (int j = 1; j < MAXLOG; j++){
        for (int i = 1; i <= n; i++){
            jumps[i][j] = jumps[jumps[i][j-1]][j-1];
            milk[i][j] = milk[i][j-1] | milk[milk[i][j-1]][j-1];
        }
    }

    for (int i = 0; i < m; i++){
        cin >> a >> b >> breed;
        int to_find = (breed == 'H') ? 1 : 2;
        cout << lca(a, b, to_find);
    }

    return 0;
}

PS. czy ktoś mógłby podesłać kod w Pythonie który generuje drzewo? (spójny graf o liczbie krawędzi o jeden mniejszych niż liczba wierzchołków bez cykli)

EDIT: Napisałem generator testów

Kopiuj
from random import randint
from random import choice

def main():
    with open("test.in", "w") as file:
        n = 10
        m = 5
        file.write(str(n) + " " + str(m) + "\n")
        for i in range(n):
            file.write(choice("HG"))
        file.write("\n")
        for i in range(2, n+1):
            file.write(str(randint(1, i-1)) +  " " + str(i) + "\n")
        for i in range(m):
            file.write(str(randint(1, n)) + " " + str(randint(1, n)) + " " + choice("HG") + "\n")
        
main()

EDIT: Okazuje się że wartość w milk nie jest indeksem poprzednika xD

Kopiuj
milk[i][j] = milk[i][j-1] | milk[milk[i][j-1]][j-1];

Trzeba zamienić na:

Kopiuj
milk[i][j] = milk[i][j-1] | milk[jumps[i][j-1]][j-1];