HackerRank Journey Scheduling problem solution

In this HackerRank Journey Scheduling problem solution, Fedya has prepared a list of M possible journeys. Each one is characterized by two integers – the starting city V and the total number of cities to be visited, K. For each of them, he is keen to know the total distance traveled by him.

Problem solution in Python.

#!/bin/python3

import math
import os
import random
import re
import sys
sys.setrecursionlimit(10**9)

#
# Complete the 'journeyScheduling' function below.
#
# The function is expected to return an INTEGER_ARRAY.
# The function accepts following parameters:
#  1. INTEGER n
#  2. 2D_INTEGER_ARRAY roads
#  3. 2D_INTEGER_ARRAY journeys
#

def journeyScheduling(n, roads, journeys):
    # Write your code here
    adj_list = [[] for i in range(n+1)]
    for u, v in roads:
        adj_list[u].append(v)
        adj_list[v].append(u)
    
    start = 1
    down_max_one = [(0, None)] * (n+1)
    down_max_two = [0] * (n+1)
    visited = [False] * (n+1)
    queue = [(start, None)]
    
    while len(queue):
        edge = queue.pop()
        u, p = edge
        if not visited[u]:
            queue.append(edge)
            visited[u] = True
            for v in adj_list[u]:
                if visited[v]:
                    continue
                # dfs_down(v)
                queue.append((v, u))
            
            continue
        
        d_one = 0
        nl_one = None
        d_two = 0
        nl_two = None
        for v in adj_list[u]:
            if v == p:
                continue
            d_v_one, nl_v_one = down_max_one[v]
            d_v_one += 1
            if d_v_one > d_one:
                d_two = d_one
                # nl_two = nl_one
                d_one = d_v_one
                nl_one = v

            elif d_v_one >= d_two:
                d_two = d_v_one
                # nl_two = v
                                    
        down_max_one[u] = (d_one, nl_one)
        down_max_two[u] = d_two
    
    # dfs_down(start)
        
    visited = [False] * (n+1)
    up_max = [0] * (n+1)
    dist_max = [0] * (n+1)
    queue = [(start, None)]
    
    while len(queue):
        edge = queue.pop()
        u, p = edge
        visited[u] = True

        if p is None:
            up_u = 0
            # up_nl_u = None#set()
        else:
            up_p = up_max[p]
            up_u_p = up_p + 1
            d_p_one, d_nl_p_one = down_max_one[p]
            d_p_two = down_max_two[p]
            
            if u != d_nl_p_one:
                d_p_v = d_p_one + 1
            else:
                d_p_v = d_p_two + 1
            
            up_u = max(up_u_p, d_p_v)
            
        up_max[u] = up_u
        d_u, d_nl_u = down_max_one[u]
        
        dist_max[u] = max(d_u, up_u)
        
        for v in adj_list[u]:
            if visited[v]:
                continue
            queue.append((v, u))


    # dfs_max_dist(start, None)
    diameter = max(dist_max)
    
    # print(diameter)
    # print(dist_max)
    m = len(journeys)
    res = [0] * m
    for i in range(m) :
        v, k = journeys[i]
        max_v = dist_max[v]
        res[i] = max_v + (k-1)*diameter
    return res
    

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    first_multiple_input = input().rstrip().split()

    n = int(first_multiple_input[0])

    m = int(first_multiple_input[1])

    roads = [[] for i in range(n-1)]

    for i in range(n - 1):
        road_inputs = input().rstrip().split()
        roads[i] = (int(road_inputs[0]), int(road_inputs[1]))

    journeys = [[] for i in range(m)]

    for j in range(m):
        journey_inputs = input().rstrip().split()
        journeys[j] = (int(journey_inputs[0]), int(journey_inputs[1]))

    result = journeyScheduling(n, roads, journeys)
    # result = [0, 11]
    fptr.write('n'.join(map(str, result)))
    fptr.write('n')

    fptr.close()

Problem solution in Java.

import java.io.*;
import java.math.*;
import java.text.*;
import java.util.*;
import java.util.regex.*;

public class Solution {
    private static class Node {
        final int index;
        final Set<Node> neighbors = new HashSet<>();

        Node(int index) {
            this.index = index;
        }
    }

    /*
     * Complete the journeyScheduling function below.
     */
    static long[] journeyScheduling(int n, int[][] roads, int[][] journeys) {
        Node[] nodes = new Node[n];
        for (int i = 0; i < n; i++) {
            nodes[i] = new Node(i);
        }
        for (int[] road : roads) {
            nodes[road[0] - 1].neighbors.add(nodes[road[1] - 1]);
            nodes[road[1] - 1].neighbors.add(nodes[road[0] - 1]);
        }

        Node start = findEnd(nodes[0]);
        Node end = findEnd(start);
        List<Node> diameterPath = findPath(start, end);
        int diameter = diameterPath.size() - 1;
        int[] distances = new int[n];
        if ((diameter & 1) == 0) {
            maxDistances(diameterPath.get(diameter/2), null, diameter/2, distances);
        } else {
            maxDistances(diameterPath.get(diameter/2), diameterPath.get(diameter/2 + 1), diameter/2 + 1, distances);
            maxDistances(diameterPath.get(diameter/2 + 1), diameterPath.get(diameter/2), diameter/2 + 1, distances);
        }
//        System.out.println(String.format("Diameter: %d, distances: %s", diameter, Arrays.toString(distances)));

        long[] results = new long[journeys.length];
        for (int i = 0; i < journeys.length; i++) {
            results[i] = distances[journeys[i][0] - 1] + ((long) diameter)*(journeys[i][1] - 1);
        }
        return results;
    }

    private static class State {
        final Node cur;
        final Node prev;
        final Iterator<Node> children;
        final int distance;

        State(Node cur, Node prev, int distance) {
            this.cur = cur;
            this.prev = prev;
            this.children = cur.neighbors.iterator();
            this.distance = distance;
        }
    }

    private static Node findEnd(Node cur) {
        Node end = cur;
        int maxDistance = -1;
        Deque<State> stack = new ArrayDeque<>();
        stack.addFirst(new State(cur, null, 0));
        while (!stack.isEmpty()) {
            State state = stack.peekFirst();
            if (state.children.hasNext()) {
                Node child = state.children.next();
                if (child == state.prev) {
                    // Do nothing
                } else if (child.neighbors.size() == 1) {
                    if (state.distance > maxDistance) {
                        maxDistance = state.distance;
                        end = child;
                    }
                } else {
                    stack.addFirst(new State(child, state.cur, state.distance + 1));
                }
            } else {
                stack.removeFirst();
            }
        }
        return end;
    }

    private static List<Node> findPath(Node cur, Node goal) {
        Deque<State> stack = new ArrayDeque<>();
        stack.addFirst(new State(cur, null, 0));
        while (stack.peekFirst().cur != goal) {
            State state = stack.peekFirst();
            if (state.children.hasNext()) {
                Node child = state.children.next();
                if (child != state.prev) {
                    stack.addFirst(new State(child, state.cur, 0));
                }
            } else {
                stack.removeFirst();
            }
        }

        List<Node> path = new ArrayList<>();
        while (!stack.isEmpty()) {
            path.add(stack.removeFirst().cur);
        }
        return path;
    }

    private static void maxDistances(Node cur, Node prev, int distance, int[] distances) {
        distances[cur.index] = distance;
        Deque<State> stack = new ArrayDeque<>();
        stack.addFirst(new State(cur, prev, distance));
        while (!stack.isEmpty()) {
            State state = stack.peekFirst();
            if (state.children.hasNext()) {
                Node child = state.children.next();
                if (child != state.prev) {
                    distances[child.index] = state.distance + 1;
                    stack.addFirst(new State(child, state.cur, state.distance + 1));
                }
            } else {
                stack.removeFirst();
            }
        }
    }

    private static final Scanner scanner = new Scanner(System.in);

    public static void main(String[] args) throws IOException {
        BufferedWriter bufferedWriter = new BufferedWriter(new FileWriter(System.getenv("OUTPUT_PATH")));

        String[] nm = scanner.nextLine().split(" ");
        scanner.skip("(rn|[nru2028u2029u0085])*");

        int n = Integer.parseInt(nm[0]);

        int m = Integer.parseInt(nm[1]);

        int[][] roads = new int[n-1][2];

        for (int roadsRowItr = 0; roadsRowItr < n-1; roadsRowItr++) {
            String[] roadsRowItems = scanner.nextLine().split(" ");
            scanner.skip("(rn|[nru2028u2029u0085])*");

            for (int roadsColumnItr = 0; roadsColumnItr < 2; roadsColumnItr++) {
                int roadsItem = Integer.parseInt(roadsRowItems[roadsColumnItr]);
                roads[roadsRowItr][roadsColumnItr] = roadsItem;
            }
        }

        int[][] journeys = new int[m][2];

        for (int journeysRowItr = 0; journeysRowItr < m; journeysRowItr++) {
            String[] journeysRowItems = scanner.nextLine().split(" ");
            scanner.skip("(rn|[nru2028u2029u0085])*");

            for (int journeysColumnItr = 0; journeysColumnItr < 2; journeysColumnItr++) {
                int journeysItem = Integer.parseInt(journeysRowItems[journeysColumnItr]);
                journeys[journeysRowItr][journeysColumnItr] = journeysItem;
            }
        }

        long[] result;
//        try {
            result = journeyScheduling(n, roads, journeys);
/*        } catch (StackOverflowError e) {
            result = new long[1];
        }*/

        for (int resultItr = 0; resultItr < result.length; resultItr++) {
            bufferedWriter.write(String.valueOf(result[resultItr]));

            if (resultItr != result.length - 1) {
                bufferedWriter.write("n");
            }
        }

        bufferedWriter.newLine();

        bufferedWriter.close();

        scanner.close();
    }
}

Problem solution in C++.

#include <cmath>
#include <cstdio>
#include <vector>
#include <iostream>
#include <algorithm>
#include <map>

using namespace std;

vector<int> visit;
int len;

int dfs(vector<map<int, int> >& g, int node, int dist) {
    //cout << "dfs " << node << " dist " << dist << endl;
    visit[node] = 1;
    int max1 = dist, max2 = 0, max_dep = 0;
    for(auto iter = g[node].begin(); iter != g[node].end(); ++iter) {
        if(iter->second >= max1) {
            max2 = max1;
            max1 = iter->second;
        } else {
            max2 = max(iter->second, max2);
        }
    }
    //cout << "nod " << node << " max " << max1 << " " << max2 << endl;
    for(auto iter = g[node].begin(); iter != g[node].end(); ++iter) {
        if(visit[iter->first]) {
            iter->second = dist;
        } else {
            iter->second = dfs(g, iter->first, 1 + (iter->second != max1 ? max1 : max2));
            max_dep = max(max_dep, iter->second);
        }
    }
    for(auto it2 = g[node].begin(); it2 != g[node].end(); ++it2) {
        //cout << "(" << it2->first << ", " <<it2->second << ") ";
    }
    //cout << endl << len << endl;
    len = max(len, dist);
    return max_dep + 1;
}

int main() {
    int n, m;
    cin >> n >> m;
    vector<map<int, int> > g(n);
    for(int i = 0; i < n - 1; ++i) {
        int a, b;
        cin >> a >> b;
        g[a - 1][b - 1] = 0;
        g[b - 1][a - 1] = 0;
    }
    visit.assign(n, 0);
    dfs(g, 0, 0);
    visit.assign(n, 0);
    dfs(g, 0, 0);
    /*
    for(auto it1 = g.begin(); it1 != g.end(); ++it1) {
        for(auto it2 = (*it1).begin(); it2 != (*it1).end(); ++it2) {
            cout << "(" << it2->first << ", " <<it2->second << ") ";
        }
        cout << endl;
    }*/
    
    //cout << "len " << len << endl;
    
    while(m--) {
        int v, dist = 0;
        long long k, sum = 0;
        cin >> v >> k;
        for(auto iter = g[v - 1].begin(); iter != g[v - 1].end(); ++iter) {
            dist = max(dist, iter->second);
        }
        sum = dist + (k - 1) * len;
        cout << sum << endl;
    }
    return 0;
}

Problem solution in C.

#include <stdio.h>
        #include <stdlib.h>
        #include <string.h>
        typedef struct _node{
          int x;
          int w;
          struct _node *next;
        } lnode;
        void insert_edge(int x,int y,int w);
        void dfs1(int x);
void dfs2(int x,int y);
int max(int x,int y);
        int h[100000],sh[100000],l[100000],u[100000],trace[100000];
        lnode *table[100000]={0};
         
        int main(){
          int N,M,x,y,i;
            scanf("%d%d",&N,&M);
            for(i=0;i<N-1;i++){
              scanf("%d%d",&x,&y);
              insert_edge(x-1,y-1,1);
            }
            memset(trace,0,sizeof(trace));
            dfs1(0);
            memset(trace,0,sizeof(trace));
            dfs2(0,-1);
            for(i=0;i<M;i++){
              scanf("%d%d",&x,&y);
              printf("%lldn",max(h[x-1],u[x-1])+(y-1)*(long long)l[0]);
            }
          return 0;
        }
        void insert_edge(int x,int y,int w){
          lnode *t=malloc(sizeof(lnode));
          t->x=y;
          t->w=w;
          t->next=table[x];
          table[x]=t;
          t=malloc(sizeof(lnode));
          t->x=x;
          t->w=w;
          t->next=table[y];
          table[y]=t;
          return;
        }
        void dfs1(int x){
          lnode *p;
          trace[x]=1;
          h[x]=sh[x]=l[x]=0;
          for(p=table[x];p;p=p->next)
            if(!trace[p->x]){
              dfs1(p->x);
              if(h[p->x]+p->w>=h[x]){
                sh[x]=h[x];
                h[x]=h[p->x]+p->w;
              }
              else if(h[p->x]+p->w>sh[x])
                sh[x]=h[p->x]+p->w;
              if(l[p->x]>l[x])
                l[x]=l[p->x];
            }
          if(h[x]+sh[x]>l[x])
            l[x]=h[x]+sh[x];
          return;
        }
        void dfs2(int x,int y){
          lnode *p;
          trace[x]=1;
if(y==-1)
u[x]=0;
else
if(h[y]==h[x]+1)
u[x]=max(u[y]+1,sh[y]+1);
else
u[x]=max(u[y]+1,h[y]+1);
          for(p=table[x];p;p=p->next)
            if(!trace[p->x])
              dfs2(p->x,x);
          return;
        }
int max(int x,int y){
  return (x>y)?x:y;
}