HackerRank Jeanie’s Route problem solution

In this HackerRank Jeanie’s Route problem solution Byteland has N cities (numbered from 1 to N) and N – 1 bidirectional road. It is guaranteed that there is a route from any city to any other city.

Jeanie is a postal worker who must deliver K letters to various cities in Byteland. She can start and end her delivery route in any city. Given the destination cities for K letters and the definition of each road in Byteland, find and print the minimum distance Jeanie must travel to deliver all K letters.

Problem solution in Python.

def farthest_node_from_start(nodes,is_goal,subtree,starting_node):
    visited=[False for _ in nodes]
    st=[(0,starting_node)]
    visited[starting_node]=True
    largest = 0,starting_node
    all_dist_sum = 0
    while st:
        dist,u=st.pop()
        for v,d in nodes[u]:
            if not visited[v] and subtree[v]:
                elm=dist+d,v
                all_dist_sum+=d
                if is_goal[v]:
                    largest=max(largest,elm)
                st.append(elm)
                visited[v]=True
    return all_dist_sum,largest

def make_subtree(nodes,is_goal,starting_node):
    subtree=is_goal[:]
    for u,v in dfs(nodes, starting_node):
        subtree[u]=subtree[u] or subtree[v]
    return subtree

def dfs(nodes, starting_node):
    stack = [(starting_node, v) for v,d in nodes[starting_node]]
    edges = list()
    while stack:
        u, v = stack.pop()
        edges.append((u, v))
        stack.extend((v, w) for w,d in nodes[v] if w != u)
    edges.reverse()
    return edges

N, K = list(map(int,input().strip().split()))
goals = list(map(int,input().strip().split()))
is_goal=[False]*(N+1)

for item in goals:
    is_goal[item]=True
    
nodes=[[] for _ in range(N+1)]

for _ in range(N-1):
    u,v,d = list(map(int,input().strip().split()))
    nodes[u].append((v,d))
    nodes[v].append((u,d))

subtree = make_subtree(nodes,is_goal,goals[0])
a,(b,c)=farthest_node_from_start(nodes,is_goal,subtree,goals[0])
Distance,(d,_)=farthest_node_from_start(nodes,is_goal,subtree,c)
print(2*Distance-d)

Problem solution in Java.

import java.util.Scanner;
import java.util.Queue;
import java.util.ArrayDeque;

class Solution{
    static boolean[] prune(int[][][] adj, boolean[] isLett){
    int n=adj.length;
    int[] degree=new int[n];
    for(int i=0;i<n;++i) degree[i]=adj[i].length;
    boolean[] rem=new boolean[n];
    Queue<Integer> q=new ArrayDeque<>();
    for(int i=0;i<n;++i){
        if(!isLett[i] && degree[i]==1) q.add(i);
    }
    while(!q.isEmpty()){
        int leaf=q.remove();
        rem[leaf]=true;
        for(int[] edge: adj[leaf]){
        int node=edge[0];
        if(isLett[node]) break;
        else if(!rem[node]){
            if(--degree[node] == 1){
            q.add(node);
            break;
            }
        }
        }
    }
    return rem;
    }
    static int[] bfs(int[][][] adj, boolean[] rem, int source){
    int n=adj.length, unvis=-1;
    int[] dist=new int[n];
    for(int i=0;i<n;++i) dist[i]=unvis;
    Queue<Integer> q=new ArrayDeque<>();
    q.add(source);
    dist[source]=0;
    int best=0, total=0;
    while(!q.isEmpty()){
        int x=q.remove();
        for(int[] edge: adj[x]){
        int to=edge[0];
        if(!rem[to] && dist[to]==unvis){
            int weight=edge[1];
            total+=weight;
            q.add(to);
            dist[to]=dist[x]+weight;
            if(dist[to]>dist[best]) best=to;
        }
        }
    }
    int[] result={total,dist[best],best};
    return result;
    }
    static int solve(int[][][] adj, int[] lett){
    boolean[] isLett=new boolean[adj.length];
    for(int i: lett) isLett[i]=true;    
    boolean[] rem=prune(adj,isLett);
    int[] result=bfs(adj,rem,lett[0]);
    int totalWeight=result[0], sink=result[2];
    result=bfs(adj,rem,sink);
    int diameter=result[1];
    return 2*totalWeight-diameter;
    }
    static int[][][] weightedAdjacency(int n, int[] from, int[] to, int[] d){
    int[] count=new int[n];
    for(int f: from) ++count[f];
    for(int t: to) ++count[t];
    int[][][] adj=new int[n][][];
    for(int i=0;i<n;++i) adj[i]=new int[count[i]][];
    for(int i=0;i<from.length;++i){
        adj[from[i]][--count[from[i]]]=new int[]{to[i],d[i]};
        adj[to[i]][--count[to[i]]]=new int[]{from[i],d[i]};
    }
    return adj;
    }
    public static void main(String[] args){
    Scanner sc=new Scanner(System.in);
    int n=sc.nextInt(), k=sc.nextInt();
    int[] lett=new int[k];
    for(int i=0;i<k;++i) lett[i]=sc.nextInt()-1;
    int[] from=new int[n-1], to=new int[n-1], d=new int[n-1];
    for(int i=0;i<n-1;++i){
        from[i]=sc.nextInt()-1;
        to[i]=sc.nextInt()-1;
        d[i]=sc.nextInt();
    }
    sc.close();
    int[][][] adj=weightedAdjacency(n,from,to,d);
    System.out.println(solve(adj,lett));
    }
}

Problem solution in C++.

#include <bits/stdc++.h>
#define pb push_back
#define mp make_pair
#define MA(x,y) ((x)>(y)?(x):(y))
using namespace std;
const int N=500002;
const int inf=1000000000;
int n,K,s[N],SUM,M;
bool f[N];
vector <int> v[N],d[N];

int input(){
    scanf("%d%d",&n,&K);
    for (int i=0,x;i<K;i++){
        scanf("%d",&x);
        s[x]=1;
        f[x]=1;
    }
    for (int i=1,x,y,z;i<n;i++){
        scanf("%d%d%d",&x,&y,&z);
        v[x].pb(y);
        v[y].pb(x);
        d[x].pb(z);
        d[y].pb(z);
    }
    return 0;
}

int go(int x,int from){
    int D1=-inf;
    int D2=-inf;

    for (int i=0;i<v[x].size();i++)
    if (v[x][i]!=from){
        D1=max(D1,go(v[x][i],x)+d[x][i]);
        if (D1>D2) swap(D1,D2);
        s[x]+=s[v[x][i]];
        if (0<s[v[x][i]] && s[v[x][i]]<K) SUM+=d[x][i];
    }

    if (D1>0) M=max(M,D1+D2);

    if (D2>0 && f[x]) M=max(M,D2);

    if (f[x]) D2=max(D2,0);

    return D2;
}

int sol(){
    go(1,1);
    printf("%dn",SUM*2-M);
    return 0;
}

int main() {
    input();
    sol();
    return 0;
}

Problem solution in C.

#include <assert.h>
#include <limits.h>
#include <math.h>
#include <stdbool.h>
#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 z);
int max(int x,int y);
int mark[100000]={0},count[100000],trace[100000],l[100000],ul[100000],s[100000],sc[100000],len[100000],ulen[100000],ans=2147483647;
lnode *table[100000]={0};

int main(){
  int N,K,x,y,w,i;
  scanf("%d%d",&N,&K);
  while(K--){
    scanf("%d",&x);
    mark[x-1]=1;
  }
  for(i=0;i<N-1;i++){
    scanf("%d%d%d",&x,&y,&w);
    insert_edge(x-1,y-1,w);
  }
  memset(trace,0,sizeof(trace));
  dfs1(0);
  memset(trace,0,sizeof(trace));
  dfs2(0,-1,-1);
  printf("%d",ans);
  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;
  if(mark[x])
    count[x]=1;
  else
    count[x]=0;
  l[x]=s[x]=sc[x]=len[x]=0;
  for(p=table[x];p;p=p->next)
    if(!trace[p->x]){
      dfs1(p->x);
      count[x]+=count[p->x];
      if(count[p->x]){
        sc[x]++;
        len[x]+=len[p->x]+2*p->w;
        if(l[p->x]+p->w>l[x]){
          s[x]=l[x];
          l[x]=l[p->x]+p->w;
        }
        else if(l[p->x]+p->w>s[x])
          s[x]=l[p->x]+p->w;
      }
    }
  return;
}
void dfs2(int x,int y,int z){
  lnode *p;
  int rl,rs;
  trace[x]=1;
  if(y==-1)
    ulen[x]=ul[x]=0;
  else{
    if(count[0]-count[x]){
      ulen[x]=len[y]+ulen[y]-len[x];
      if(!count[x])
        ulen[x]+=2*z;
      if(l[y]==z+l[x])
        ul[x]=z+max(s[y],ul[y]);
      else
        ul[x]=z+max(l[y],ul[y]);
    }
    else
      ulen[x]=ul[x]=0;
  }
    rl=l[x];
    rs=s[x];
    if(ul[x]>rl){
      rs=rl;
      rl=ul[x];
    }
    else if(ul[x]>rs)
      rs=ul[x];
    if(len[x]+ulen[x]-rl-rs<ans)
      ans=len[x]+ulen[x]-rl-rs;
  for(p=table[x];p;p=p->next)
    if(!trace[p->x])
      dfs2(p->x,x,p->w);
  return;
}
int max(int x,int y){
  return (x>y)?x:y;
}