In this HackerRank Matrix Interview preparation kit problem a Heap there is Given a list of edges and times, determine the minimum time to stop the attack.


HackerRank Matrix Interview preparation kit solution


Problem solution in Python programming.

class Edge(object):
    def __init__(self, city1: int, city2: int, weight: int):
        self.cities = [city1, city2]
        self.weight = weight

[N, K] = [int(i) for i in input().split(' ')]
edges = list()
for i in range(N - 1):
    line = input()
    #print(line)
    [city1, city2, weight] = [int(j) for j in line.split(' ')]
    edges.append(Edge(city1, city2, weight))

edges.sort(key = lambda x: x.weight, reverse = True)

cityToGroup = dict()
machineCities = dict()
for i in range(K):
    mc = int(input())
    machineCities[mc] = True
    cityToGroup[mc] = mc
    
result = 0
groups = dict()
for i in range(N):
    groups[i] = [i]
    cityToGroup[i] = i

for e in edges:
    city1 = e.cities[0]
    city2 = e.cities[1]
    g1 = cityToGroup[city1]
    g2 = cityToGroup[city2]
    if g1 == g2:
        continue
    if machineCities.get(g1, None) is not None: # in this group we have a machine
        if machineCities.get(g2, None) is not None: # do not add this edge, as it connects the machines
            result += e.weight
        else:
            for c in groups[g2]:
                cityToGroup[c] = g1
                groups[g1].append(c)
            del groups[g2]
    else: # Either the second group has a machine or not, we add all g1 cities to it and delete the first g1
        for c in groups[g1]:
            groups[g2].append(c)
            cityToGroup[c] = g2
        del groups[g1]
print(result)


Problem solution in Java Programming.

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

public class Solution {

    // Complete the minTime function below.
    static int minTime(int[][] roads, int[] machines) {
    int n = roads.length + 1;
    int[] sum_edge = new int[n];
    int[] max_edge = new int[n];
    int[] processed_connections = java.util.stream.IntStream.generate(() -> 2).limit(n).toArray();
    boolean[] is_visited = new boolean[n];
    boolean[] is_machine = new boolean[n];
    Arrays.stream(machines).forEach(m -> is_machine[m]=true);
    List<int[]>[] edge_list = java.util.stream.Stream.generate(ArrayList::new).limit(n).toArray(List[]::new);
    for(int[] road: roads) {
        edge_list[road[0]].add(road);
        edge_list[road[1]].add(road);
    }

    int result = 0;
    for(int i = 0; i < n; i++) {
        if(edge_list[i].size()!=1 || is_visited[i]) continue;

        is_visited[i] = true;
        boolean found = is_machine[i];
        int curr_node = i;
        int[] next_edge = edge_list[curr_node].get(0);
        int min_edge_time = next_edge[2];
        for(boolean finished = false; !finished;) {
            curr_node = (curr_node==next_edge[0]? next_edge[1] : next_edge[0]);
            int size = edge_list[curr_node].size();
            if(size == 1) {
                if(found && is_machine[curr_node]) result+=min_edge_time;
                is_visited[curr_node] = true;
                finished = true;
            } else if(size > 2) {
                if(is_machine[curr_node]) {
                    if(found) result+=min_edge_time;
                    found = true;
                } else if(found) {
                    sum_edge[curr_node] += min_edge_time;
                    max_edge[curr_node] = Math.max(max_edge[curr_node], min_edge_time);
                }
        next_edge[2] = 0;
                if(++processed_connections[curr_node]<=size) finished=true;
                else {
        boolean has_more = false;
        for(int[] edge : edge_list[curr_node]) {
            if(edge[2] != 0) {
                next_edge = edge;
                has_more = true;
                break;
            }
        }

        if(has_more) {
                    if(is_machine[curr_node]) min_edge_time=next_edge[2];
                    else {
                        result += sum_edge[curr_node] - max_edge[curr_node];
                        found = true;
                        min_edge_time = Math.min(max_edge[curr_node], next_edge[2]);
                    }
        } else finished = true;
                }
            } else {
                next_edge = (next_edge==edge_list[curr_node].get(0)? edge_list[curr_node].get(1) : edge_list[curr_node].get(0));
                if(!is_machine[curr_node]) min_edge_time=Math.min(min_edge_time, next_edge[2]);
                else {
                    if(found) result+=min_edge_time;
                    found = true;
                    min_edge_time = next_edge[2];
                }
            }
        }
    }
    return result;
}

    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[] nk = scanner.nextLine().split(" ");

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

        int k = Integer.parseInt(nk[1]);

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

        for (int i = 0; i < n - 1; i++) {
            String[] roadsRowItems = scanner.nextLine().split(" ");
            scanner.skip("(\r\n|[\n\r\u2028\u2029\u0085])?");

            for (int j = 0; j < 3; j++) {
                int roadsItem = Integer.parseInt(roadsRowItems[j]);
                roads[i][j] = roadsItem;
            }
        }

        int[] machines = new int[k];

        for (int i = 0; i < k; i++) {
            int machinesItem = scanner.nextInt();
            scanner.skip("(\r\n|[\n\r\u2028\u2029\u0085])?");
            machines[i] = machinesItem;
        }

        int result = minTime(roads, machines);

        bufferedWriter.write(String.valueOf(result));
        bufferedWriter.newLine();

        bufferedWriter.close();

        scanner.close();
    }
}


Problem solution in C++ programming.

#include <cstdio>
#include <vector>

using namespace std;

#define forn(i,n) for (int i = 0; i < (n); i++)
typedef long long LL;
typedef pair<int, int> ii;

#define N 100004
const LL inf = 1000000000000000000LL;
int n, m;
bool marked[N];
LL f[N], g[N];
vector<ii> a[N];

void go(int v, int parent) {
	LL sum = 0, mx = -inf;
	for (vector<ii>::iterator it = a[v].begin(); it != a[v].end(); ++it) {
		int x = it->first;
		if (x == parent) continue;
		go(x, v);
		if (f[x] == inf) continue;
		LL u = min(g[x], f[x] + it->second);
		sum += u;
		mx = max(mx, u - f[x]);
	}
	if (marked[v]) f[v] = sum, g[v] = inf;
		else f[v] = sum - mx, g[v] = sum;
}

int main() {
	scanf("%d%d", &n, &m);
	forn(_, n-1) {
		int x, y, z;
		scanf("%d%d%d", &x, &y, &z);
		a[x].push_back(ii(y, z));
		a[y].push_back(ii(x, z));
	}
	forn(_, m) {
		int x;
		scanf("%d", &x);
		marked[x] = 1;
	}
	go(0, -1);
	printf("%lld\n", min(f[0], g[0]));
	return 0;
}


Problem solution in C programming.

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
    union {
        struct {unsigned city_x, city_y;};
        unsigned long packd;
    };
    unsigned cost;
} road_t;

typedef struct {
    unsigned *indices, terminal;
    road_t *items;
    unsigned count;
} min_heap_t;


static inline road_t *min_heap_pop(min_heap_t *self) {
    road_t *smallest = &self->items[self->indices[1]];

    unsigned
        location = self->indices[self->count],
        min_cost = self->items[location].cost,
        parent = 1,
        child;

    self->indices[self->count--] = self->terminal;

    for (child = 2; child <= self->count; child <<= 1) {
        child |= self->items[self->indices[child | 1U]].cost < self->items[self->indices[child]].cost;

        if (self->items[self->indices[child]].cost <= min_cost) {
            self->indices[parent] = self->indices[child];
            parent = child;
        } else
            break ;
    }
    self->indices[parent] = location;
    return smallest;
}

static inline void min_heap_push(min_heap_t *self, unsigned at) {
    unsigned root;
    for (root = ++self->count; root > 1 && self->items[at].cost < self->items[self->indices[root >> 1]].cost; root >>= 1)
        self->indices[root] = self->indices[root >> 1];
    self->indices[root] = at;
}

#define have_seen(self, id)  ((self)[(id) >> 5]  & (1U << ((id) & 31U)))
#define inv_seen(self, id)   ((self)[(id) >> 5] ^= (1U << ((id) & 31U)))

unsigned
    occupied[3125] = {[0 ... 3124] = 0},
    set[3125] = {[0 ... 3124] = 0},
    stack[100000],
    history[100000];

int main() {
    unsigned at, city_cnt, edge_cnt, machine_cnt;
    scanf("%u %u", &city_cnt, &machine_cnt);

    unsigned long indices[city_cnt + 1];
    memset(indices, 0, sizeof(indices));

    edge_cnt = city_cnt - 1;
    road_t *road = malloc(sizeof(*road) * (edge_cnt << 1));
    for (at = 0; at < edge_cnt; at++) {
        scanf("%u %u %u", &road[at].city_x, &road[at].city_y, &road[at].cost);
        road[at + edge_cnt].cost = road[at].cost;
        road[at + edge_cnt].packd = (road[at].packd << 32) | (road[at].packd >> 32);
        indices[road[at].city_x]++;
        indices[road[at].city_y]++;
    }

    edge_cnt <<= 1;
    road_t ordered[edge_cnt | 1];
    for (at = 0; ++at < city_cnt; indices[at] += indices[at - 1]);
    for (at = edge_cnt; at--; ordered[--indices[road[at].city_x]] = road[at]);
    
    free(road);

    indices[city_cnt] = edge_cnt;
    ordered[edge_cnt].packd = ((unsigned long)city_cnt << 32) | city_cnt;
    ordered[edge_cnt].cost = 0xFFFFFFFFU;

    for (; machine_cnt--; inv_seen(occupied, at))
        scanf("%u", &at);


    // take all the edges and insert them into priority queue ordered by increasing weight
    // while the queue is not empty pop and temporarily remove it,
    // do a dfs on both cities,
    //  if they both lead to an occupied city, permanently remove the edge,
    //  otherwise restore the edge.

    unsigned buffer[edge_cnt | 1];
    min_heap_t edges = {.indices = buffer, .items = ordered, .count = 0, .terminal = edge_cnt};
    
    for (at = edge_cnt | 1; at--; edges.indices[at] = edge_cnt);
    for (at = edge_cnt; at--; min_heap_push(&edges, at));

    unsigned length, city;
    unsigned long total = 0;

    road_t *neighbors;
    while (edges.count) {
        road = min_heap_pop(&edges);
        if (road->city_x == road->city_y) // edge has being removed ...
            continue;

        unsigned
            city_x = road->city_x,
            city_y = road->city_y;

        length = 0;
        stack[0] = city_x;
        road->city_y = city_x;
        for (at = 1; at && have_seen(occupied, stack[at - 1]) == 0; ) {
            city = stack[--at];
            inv_seen(set, city);
            history[length++] = city;
            for (neighbors = &ordered[indices[city]]; neighbors->city_x == city; neighbors++)
                if (have_seen(set, neighbors->city_y) == 0)
                    stack[at++] = neighbors->city_y;
        }

        if (at) { // city_x lead to an occupied city check city_y;
            road_t *other_road;
            for (other_road = &ordered[indices[city_y]]; other_road->city_y != city_x; other_road++);
            other_road->city_y = city_y;

            stack[0] = city_y;
            for (at = 1; at && have_seen(occupied, stack[at - 1]) == 0; ) {
                city = stack[--at];
                inv_seen(set, city);
                history[length++] = city;
                for (neighbors = &ordered[indices[city]]; neighbors->city_x == city; neighbors++)
                    if (have_seen(set, neighbors->city_y) == 0)
                        stack[at++] = neighbors->city_y;
            }

            if (at == 0) // restore edge, since we didn't reach an occupied city.
                other_road->city_y = city_x;
        }

        if (at)
            total += road->cost;
        else
            road->city_y = city_y;

        for (; length--; inv_seen(set, history[length]));
    }

    printf("%lu", total);
    return 0;
}


Problem solution in JavaScript programming.

'use strict';

const fs = require('fs');

process.stdin.resume();
process.stdin.setEncoding('utf-8');

let inputString = '';
let currentLine = 0;

process.stdin.on('data', inputStdin => {
    inputString += inputStdin;
});

process.stdin.on('end', function() {
    inputString = inputString.replace(/\s*$/, '')
        .split('\n')
        .map(str => str.replace(/\s*$/, ''));

    main();
});

function readLine() {
    return inputString[currentLine++];
}

// Complete the minTime function below.
function minTime(roads, machines) {
    machines = new Set(machines);
    
    let edges = Array(roads.length + 1);
    for(let i = 0; i < edges.length; i++) {
        edges[i] = {};
    }
    
    for(let i = 0; i < roads.length; i++) {
        let f = roads[i][0];
        let t = roads[i][1];
        let cost = roads[i][2];
        edges[f][t] = cost;
        edges[t][f] = cost;
    }
    
    let machineList = Array.from(machines);
    let cost = 0;
    let deleteCount = 0;
    for(let i = 0; i < machineList.length; i++) {
        let removals = findMinimumEdge(machines, edges, machineList[i]);
        console.log(removals);
        for(let j = 0; j < removals.length; j++) {
            let removal = removals[j];
            let t = removal[0];
            let f = removal[1];
            if(!edges[t][f]) {
                continue;
            }
            cost += edges[t][f];
            delete edges[t][f];
            delete edges[f][t];
            console.log(cost);
        }
        console.log(cost);
    }
    return cost;
}

function findMinimumEdge(machines, edges, start, path, visited) {
    if(!visited) {
        visited = {};
    }
    if(!path) {
        path = [start];
    }
    
    visited[start] = true;
    let minEdges = [];
    for(let edge in edges[start]) {
        edge = parseInt(edge);
        if(visited[edge]) {
            continue;
        }
        let newPath = path.slice(0);
        newPath.push(edge);
        if(machines.has(edge)) {
            newPath.push(edge);
            let min = edges[newPath[0]][newPath[1]];
            let minEdge = [newPath[0],newPath[1]];
            for(let i = 1; i < newPath.length - 1; i++) {
                if(edges[newPath[i]][newPath[i+1]] < min) {
                    min = edges[newPath[i]][newPath[i+1]];
                    minEdge = [newPath[i],newPath[i + 1]];
                }
            }
            minEdges.push(minEdge);
            continue;
        }
        let result = findMinimumEdge(machines, edges, edge, newPath, visited);
        minEdges.push(...result);
    }
    return minEdges;
}

function main() {
    const ws = fs.createWriteStream(process.env.OUTPUT_PATH);

    const nk = readLine().split(' ');

    const n = parseInt(nk[0], 10);

    const k = parseInt(nk[1], 10);

    let roads = Array(n - 1);

    for (let i = 0; i < n - 1; i++) {
        let line = readLine();
        if(line.indexOf('green') != -1) {
            i--;
            continue;
        }
        roads[i] = line.split(' ').map(roadsTemp => parseInt(roadsTemp, 10));
    }

    let machines = [];

    for (let i = 0; i < k; i++) {
        const machinesItem = parseInt(readLine(), 10);
        machines.push(machinesItem);
    }

    const result = minTime(roads, machines);

    ws.write(result + '\n');

    ws.end();
}