HackerRank Tree Flow problem solution

In this HackerRank Tree Flow problem solution, we have an undirected connected acyclic graph and we have a weighted tree with n vertices. we need to calculate and print the maximum total value of A for a given tree T.

Problem solution in Python.

from collections import defaultdict
from sys import stdin, stdout
from heapq import heappop, heappush

def dijkstra(n, graph, u):
    distance = [float("inf")] * (n + 1)
    distance[u] = 0
    visited = [False] * (n + 1)
    visited[u] = True
    queue = [(distance[u], u)]
    while queue:
        d, u = heappop(queue)
        for v, w in graph[u]:
            if not visited[v] and distance[v] > d + w:
                visited[v] = True
                distance[v] = d + w
                heappush(queue, (distance[v], v))
    return distance[1:]

def treeFlow(n, tree):
    graph = defaultdict(list)
    for u, v, w in tree:
        graph[u].append((v, w))
        graph[v].append((u, w))
    return min(sum(dijkstra(n, graph, 1)), sum(dijkstra(n, graph, n)))

if __name__ == '__main__':
    n = int(stdin.readline())
    tree = [list(map(int, stdin.readline().rstrip().split())) for _ in range(n - 1)]
    result = treeFlow(n, tree)
    stdout.write(str(result) + 'n')

Problem solution in Java.

import java.io.*;
import java.util.*;

public class Solution {
  private static InputReader in;
  private static PrintWriter out;
  
  static class Edge {
    public int to;
    public long weight;
    public Edge(int to, int weight) {
      this.to = to;
      this.weight = weight;
    }
  }
  
  
  public static void dfs(int node, int par, long cdist) {
    dist[node] = cdist;
    for (Edge e : graph[node]) {
      if (e.to == par) continue;
      dfs(e.to, node, cdist+e.weight);
    }
  }
  
  public static void bfs(int node) {
    int front = 0, back = 0;
    Arrays.fill(vis, false);
    queue[back++] = node;
    dist[node] = 0;
    vis[node] = true;
    while (front < back) {
      int cur = queue[front++];
      for (Edge e : graph[cur]) {
        if (vis[e.to]) continue;
        vis[e.to] = true;
        dist[e.to] = dist[cur] + e.weight;
        queue[back++] = e.to;
      }
    }
  }
  
  public static ArrayList<Edge>[] graph;
  public static long[] dist;
  public static int[] queue;
  public static boolean[] vis;
  public static void main(String[] args) throws IOException {
    in = new InputReader(System.in);
    out = new PrintWriter(System.out, true);

    int n = in.nextInt();
    graph = new ArrayList[n];
    queue = new int[graph.length];
    vis = new boolean[graph.length];
    for (int i = 0; i < n; i++) graph[i] = new ArrayList<>();
    for (int i = 0; i < n-1; i++) {
      int a = in.nextInt()-1, b = in.nextInt()-1, c = in.nextInt();
      graph[a].add(new Edge(b,c));
      graph[b].add(new Edge(a,c));
    }
    
    dist = new long[n];
    bfs(0);
    long ans1 = 0;
    for (int i = 0; i < n; i++) ans1 += dist[i];
    bfs(n-1);
    long ans2 = 0;
    for (int i = 0; i < n; i++) ans2 += dist[i];
    out.println(Math.min(ans1, ans2));
    out.close();
    System.exit(0);
  }

  static class InputReader {
    public BufferedReader reader;
    public StringTokenizer tokenizer;

    public InputReader(InputStream stream) {
      reader = new BufferedReader(new InputStreamReader(stream), 32768);
      tokenizer = null;
    }

    public String next() {
      while (tokenizer == null || !tokenizer.hasMoreTokens()) {
        try {
          tokenizer = new StringTokenizer(reader.readLine());
        } catch (IOException e) {
          throw new RuntimeException(e);
        }
      }
      return tokenizer.nextToken();
    }

    public int nextInt() {
      return Integer.parseInt(next());
    }
  }


}

Problem solution in C++.

#include <iostream>
#include <cstdio>
#include <string.h>
#include <algorithm>
#include <vector>
#include <string>
#include <queue>
#include <stack>
#include <set>
#include <map>
#include <sstream>
#include <cmath>
#include <ctime>

using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<string> vs;
typedef vector< vector<int> > vvi;
typedef vector<ll> vl;
typedef vector< vector<ll> > vvl;

#define forn(i, n) for (int i = 0; i < (int)(n); i++)
#define forv(i, v) forn(i, v.size())
#define all(v) v.begin(), v.end()
#define mp make_pair
#define pb push_back

struct Edge {
    int v, w;
    Edge() {}
    Edge(int v, int w) : v(v), w(w) {}
};

vector< vector<Edge> > g;
vl d;

void dfs(int v, int p) {
    for (Edge& e : g[v]) {
        if (e.v == p) continue;
        d[e.v] = d[v] + e.w;
        dfs(e.v, v);
    }
}

int main() {
#ifdef NEREVAR_PROJECT
    freopen("input.txt", "r", stdin);
    freopen("output.txt", "w", stdout);
#endif
    int n; cin >> n;
    g = vector< vector<Edge> >(n);
    forn(i, n - 1) {
        int a, b, c;
        scanf("%d %d %d", &a, &b, &c);
        a--, b--;
        g[a].pb(Edge(b, c));
        g[b].pb(Edge(a, c));
    }
    d = vl(n, 0);
    dfs(0, -1);
    ll s0 = 0;
    forn(i, n) s0 += d[i];
    d = vl(n, 0);
    dfs(n - 1, -1);
    ll s1 = 0;
    forn(i, n) s1 += d[i];
    cout << min(s0, s1) << endl;
    return 0;
}

Problem solution in C.

#include<stdio.h>
#include<stdlib.h>
typedef struct _lnode
{
    int x;
    int w;
    struct _lnode *next;
}lnode;
lnode *table[500000];
void dfs(int x, int y, long long len, long long *ans)
{
    lnode *p;
    *ans += len;
    for( p = table[x] ; p ; p = p -> next )
    {
        if( p -> x != y )
        {
            dfs(p -> x, x, len+p -> w, ans);
        }
    }
    return;
}
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;
}
int main()
{
    int N, x, y, z, i;
    long long ans1, ans2;
    scanf("%d", &N);
    for( i = ans1 = ans2 = 0 ; i < N - 1 ; i++ )
    {
        scanf("%d%d%d", &x, &y, &z);
        insert_edge(x-1, y-1, z);
    }
    dfs(0, -1, 0, &ans1);
    dfs(N-1, -1, 0, &ans2);
    if( ans2 < ans1 )
    {
        ans1 = ans2;
    }
    printf("%lld", ans1);
    return 0;
}