HackerRank Journey Scheduling problem solution

In this HackerRank Journey Scheduling problem solution, Fedya is a seasoned traveller and is planning his trip to Treeland. Treeland is a country with an ancient road system which is in the form of a tree structure. N cities of Treeland are numbered by N positive integers: 1,2,3,…,N.

Fedya has not yet decided the starting point (city) of his journey and the cities he will visit. But there are a few things you know about Fedya’s trip:

Fedya is fond of travelling to great distances. So if he is currently located in city V, his destination will be a city which is most distant from city V.

There might be more than 1 such cities. In that case, Fedya will choose a city that was already visited as less times as possible in this journey.

There still might be more than 1 such cities. In that case, Fedya will go to the city with the smallest number.

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 travelled by him.

HackerRank Journey Scheduling 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()

Journey Scheduling 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;
}