HackerEarth Longest Path problem solution

In this HackerEarth Longest Path problem solution You are given a tree with N nodes and N - 1 edges and an integer K. Each node has a label denoted by an integer, Ai. Your task is to find the length of the longest path such that label of all the nodes in the path are divisible by K.

Length of a path is the number of edges in that path.

HackerEarth Longest Path problem solution.

#include <bits/stdc++.h>

#define ll long long
#define ull unsigned long long
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define be begin()
#define en end()
#define le length()
#define sz size()
#define all(x) (x).begin(),(x).end()
#define alli(a, n, k) (a+k),(a+n+k)
#define REP(i, a, b, k) for(__typeof(a) i = a;i < b;i += k)
#define REPI(i, a, b, k) for(__typeof(a) i = a;i > b;i -= k)
#define REPITER(it, a) for(__typeof(a.begin()) it = a.begin();it != a.end(); ++it)
#define REPEACH(i, a) for(auto& i:a)

#define y0 sdkfaslhagaklsldk
#define y1 aasdfasdfasdf
#define yn askfhwqriuperikldjk
#define j1 assdgsdgasghsf
#define tm sdfjahlfasfh
#define lr asgasgash
#define norm asdfasdgasdgsd
#define have adsgagshdshfhds

#define eps 1e-6
#define pi 3.141592653589793

using namespace std;

template<class T> inline T gcd(T a, T b) { while(b) b ^= a ^= b ^= a %= b; return a; }
template<class T> inline T mod(T x) { if(x < 0) return -x; else return x; }

typedef vector<int> VII;
typedef vector<ll> VLL;
typedef pair<int, int> PII;
typedef pair<ll, ll> PLL;
typedef pair<int, PII > PPII;
typedef vector< PII > VPII;
typedef vector< PPII > VPPI;

const int MOD = 1e9 + 7;
const int INF = 1e9;

// Template End
const int MAX = 1e5 + 5;
int a[MAX];
int ans = 0;
VII adj[MAX];


int dfs(int x, int par) {
    int mx1 = 0, mx2 = 0;
    REPEACH(y, adj[x]) {
        if (y != par) {
            int p = dfs(y, x);
            if (p > mx1) mx2 = mx1, mx1 = p;
            else if (p > mx2) mx2 = p;
        }
    }
    if (a[x] == 0) {
        ans = max(ans, mx1 + 1);
        return mx1 + 1;
    }
    else return 0;
}

int main(int argc, char* argv[]) {
    if(argc == 2 or argc == 3) freopen(argv[1], "r", stdin);
    if(argc == 3) freopen(argv[2], "w", stdout);
    ios::sync_with_stdio(false);
    int n, k, u, v;
    cin >> n >> k;
    REP(i, 1, n+1, 1) {
        cin >> a[i];
        a[i] = (a[i] % k) ? 1 : 0;
    }
    REP(i, 1, n, 1) {
        cin >> u >> v;
        adj[u].pb(v);
        adj[v].pb(u);
    }
    ans = max(ans, dfs(1, -1));
    cout << ans-1 << endl;
    return 0;
}