Travel in HackerLand

applying parallel Binary search to solve the issue
 avatar
unknown
java
3 years ago
7.1 kB
8
Indexable
import java.io.*;
import java.math.*;
import java.security.*;
import java.text.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.function.*;
import java.util.regex.*;
import java.util.stream.*;
import static java.util.stream.Collectors.joining;
import static java.util.stream.Collectors.toList;

class Result {
    static int[] L;
    static int[] R;
    static HashMap<Integer, List<Integer>> mp;

    public static int[] travel(List<Integer> buildingsColor, List<List<Integer>> roads,
            List<List<Integer>> queries) {
        // sort by roads cost so we have ordered List that we can apply BS on
        roads.sort((List<Integer> a, List<Integer> b) -> a.get(2).compareTo(b.get(2)));

        L = new int[queries.size()];
        R = new int[queries.size()];

        mp = new HashMap<Integer, List<Integer>>();

        for (int i = 0; i < queries.size(); i++) {
            // set each query L and R value for the BinarySearch
            L[i] = 0;
            R[i] = roads.size() - 1;
            int mid = (L[i] + R[i]) / 2;
            // here we store where each query BS stopped, currently all are in the same initial value mid
            mp.computeIfAbsent(mid, k -> new ArrayList<Integer>()).add(i);
        }
        return solve(buildingsColor, roads, queries);
    }

    private static int[] solve(List<Integer> buildingsColor, List<List<Integer>> roads, List<List<Integer>> queries) {
        int n = buildingsColor.size();
        int m = roads.size();
        boolean oneMore = true;
        int[] ans = new int[queries.size()];
        for (int i = 0; i < ans.length; i++) {
            ans[i] = -1;
        }

        while (oneMore) {
            oneMore = false;
            DSU dsu = new DSU(n);
            dsu.inti(n, buildingsColor);

            for (int mid = 0; mid < m; mid++) {
                // merge this road
                dsu.merge(roads.get(mid).get(0), roads.get(mid).get(1));
                int cost = roads.get(mid).get(2); // current maximum cost of this graph
                if (mp.containsKey(mid) == false) continue;
                for (int ind : mp.get(mid)) {
                    int u = queries.get(ind).get(0);
                    int v = queries.get(ind).get(1);
                    int k = queries.get(ind).get(2);
                    //System.out.println("Current mid = " + mid);
                    if(dsu.areConnected(u, v, k)) {
                        // we found an answer, so we will try to find a better one
                        R[ind] = mid - 1;
                        ans[ind] = cost;
                        if (L[ind] <= R[ind]) {
                            int newMid = (L[ind] + R[ind]) / 2;
                            //System.out.println(mid + " " + newMid + " " + ind + " " + cost);
                            mp.computeIfAbsent(newMid, key -> new ArrayList<Integer>()).add(ind);
                            oneMore = true;
                        }
                    } else {
                        // we didn't find an answer, so we will try to find higher mid value
                        L[ind] = mid + 1;
                        if (L[ind] <= R[ind]) {
                            int newMid = (L[ind] + R[ind]) / 2;
                            mp.computeIfAbsent(newMid, key -> new ArrayList<Integer>()).add(ind);
                        }
                    }
                }
                mp.remove(mid);
            }
        }
        return ans;
    }

    static class DSU {
        int[] parent;
        ArrayList<HashSet<Integer>> sz;

        public DSU(int n) {
            parent = new int[n + 1];
            sz = new ArrayList<HashSet<Integer>>();
        }

        public void inti(int n, List<Integer> color) {
            for (int i = 1; i <= n; i++) {
                parent[i] = i;
                sz.add(new HashSet<>());
                sz.get(i - 1).add(color.get(i - 1));
            }
        }

        private int getParent(int node) {
            if (node == parent[node]) return node;
            return parent[node] = getParent(parent[node]);
        }

        private int getNodeSize(int u) {
            return sz.get(u - 1).size();
        }

        // return true if there is a path between u and v with more than or equal k colored buildings
        public boolean areConnected(int u, int v, int k) {
            //System.out.println(u + " " + v + " " + getParent(u) + " " + getParent(v) + " " + getNodeSize(getParent(u)) + " " + k);
            return (getParent(u) == getParent(v)) && (getNodeSize(getParent(u)) >= k);
        }
        public void merge(int u, int v) {
            u = getParent(u); // get u parent
            v = getParent(v); // get v parent
            if (u == v) return; // u and v already connected

            // we will be merging the smallest node to the largest
            if (sz.get(u - 1).size() < sz.get(v - 1).size()) {
                int tmp = u;
                u = v;
                v = tmp;
            }

            parent[v] = u;
            sz.get(u - 1).addAll(sz.get(v - 1));
            sz.get(v - 1).clear();
        }
    }
}

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

        String[] firstMultipleInput = bufferedReader.readLine().replaceAll("\\s+$", "").split(" ");

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

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

        int q = Integer.parseInt(firstMultipleInput[2]);

        List<Integer> t = Stream.of(bufferedReader.readLine().replaceAll("\\s+$", "").split(" "))
            .map(Integer::parseInt)
            .collect(toList());

        List<List<Integer>> roads = new ArrayList<>();

        IntStream.range(0, m).forEach(i -> {
            try {
                roads.add(
                    Stream.of(bufferedReader.readLine().replaceAll("\\s+$", "").split(" "))
                        .map(Integer::parseInt)
                        .collect(toList())
                );
            } catch (IOException ex) {
                throw new RuntimeException(ex);
            }
        });

        List<List<Integer>> queries = new ArrayList<>();

        IntStream.range(0, q).forEach(i -> {
            try {
                queries.add(
                    Stream.of(bufferedReader.readLine().replaceAll("\\s+$", "").split(" "))
                        .map(Integer::parseInt)
                        .collect(toList())
                );
            } catch (IOException ex) {
                throw new RuntimeException(ex);
            }
        });

        int[] result = Result.travel(t, roads, queries);
        StringBuilder sb = new StringBuilder();
        for(int i =0;i<result.length;i++) {
            sb.append(String.valueOf(result[i])).append('\n');
        }
           bufferedWriter.write(sb.toString());
        bufferedReader.close();
        bufferedWriter.close();
    }
}
Editor is loading...