Zadanie z USACO i algorytm LCA

Question

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ć)

#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

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()

Answer 1

Trochę mi to zajęło, ale naprawiłem ten kod. Tak wygląda po zmianach:

#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 = 18;
 
vector<vector<int>> graph(MAXN), jumps(MAXN, vector<int>(MAXLOG)), milk(MAXN, vector<int>(MAXLOG));
vector<int> depht(MAXN);

int milk_node[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;
    milk[v][0] = milk_node[p] | milk_node[v];
    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;

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

    if(x == y)
        return (found & 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];
        }
    }
    
    found |= milk[x][0];
    found |= milk[y][0];

    return (found & f);
}
 
int main(){
    setIO("milkvisits");
    int n, m;
    cin >> n >> m;
 
    char breed;
    for (int i = 1; i <= n; i++){
        cin >> breed;
        milk_node[i] = (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[jumps[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;
}

Taki największy błąd to było to, że w linii 86 (w Twoim kodzie) wybierałeś zły indeks. Ponadto milk[x][0] zawiera informacje o wierzchołku x, a jump[x][0] skacze do przodka x. Lepiej będzie jeżeli i milk[x][0], i jump[x][0] wskazują na przodka x, bo wtedy łatwiej napisać pozostałe przejścia. Dlatego dodałem zmienną milk_node[x], która trzyma rodzaj mleka dla danego wierzchołka

Comments

ale ja jestem głupi xDDD

 milk[milk[i][j-1]][j-1];

No tak napewno rodzaj mleka będzie odpowiednim indeksem xD

Ale jak dałeś że milk[i][0] przechowuje wartość wierzchołka i poprzednika i y będzie LCA x, to czy wtedy nie będziesz brał też pod uwagę poprzednika y który nie jest na ścieżce?

 

---

Linijka 51 powinna temu zaradzić. Bo jeżeli x jest potomkiem y, to wtedy found zawiera wszystkie wierzchołki od x do y włącznie