4
\$\begingroup\$

I'm trying to solve P1243E in an efficient manner. The problem in simple words is:

Given \$k\$ boxes, \$i\$-th box with \$n_i\$ numbers. All numbers are distinct. We need to select one number from each box and after permuting them in such a way that after putting exactly one number back in one box will lead to all boxes having same sum.

Input: \$k\$, then each line with \$n_i\$ and box-\$i\$ elements.

Output: Yes/No - whether such a rearrangement is possible, followed by each line containing the value to choose and the box-id it will go.

My solution is this:

First check if average is a whole number which will be the sum of each box. Then since each number is distinct, we can draw an edge from a number \$a\$ in box \$i\$ to another number \$b\$ in another box \$j\$ if we put \$a\$ in box \$j\$ and remove \$b\$ then that box's sum becomes equal to average. For this we will find the difference of the current sum in box \$j\$ and the average sum, if this is equal to \$b-a\$ then this will bring the box \$j\$'s sum to average. A permutation will be collection of cycles with distinct box ids covering all box ids. So we can just do DFS with backtracking to find such a collection of cycles.

The editorial proposes an even better method:

We can extract cycles from the graph and then convert cycle to a set with elements the box ids of covered numbers. We can then perform set cover for \$\{1,2,...k\}\$ using these sets with a dynamic programming method.

So I came up with this code which is quite faithful to the editorial that takes \$\sim 9\$ seconds for execution whereas actual time limits are \$1\$ sec. What time performance improvements can be done: (Just to confirm, there are solutions in Java by others with same logic but \$\sim 200\$ ms execution time)

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

public class P1243E {
    public static void main(String[] args) {
        InputStream inputStream = System.in;
        OutputStream outputStream = System.out;
        InputReader in = new InputReader(inputStream);
        PrintWriter out = new PrintWriter(outputStream);
        Task solver = new Task();
        boolean debug = !Boolean.parseBoolean(System.getProperty("ONLINE_JUDGE"));
        solver.solve(in, out, debug);
        out.close();
    }

    static class Task {
        public void solve(InputReader in, PrintWriter out, boolean debug) {
            int k = in.nextInt();
            int[][] boxes = new int[k][];
            HashMap<Integer, Node> nodes = new HashMap<>();
            long totSum = 0;
            long[] boxSums = new long[k];
            for (int i = 0; i < k; i++) {
                int n = in.nextInt();
                int[] box = new int[n];
                for (int j = 0; j < n; j++) {
                    int a = in.nextInt();
                    totSum += a;
                    boxSums[i] += a;
                    box[j] = a;
                    nodes.put(a, new Node(a, i));
                }
                boxes[i] = box;
            }
            if (totSum % k != 0) {
                out.println("No");
            } else {
                /* Calculate edges */
                long boxSum = totSum / k;
                long[] delta = new long[k];
                for (int i = 0; i < k; i++) {
                    delta[i] = boxSum - boxSums[i];
                }
                HashSet<Node> remaining = new HashSet<>();
                for (int i = 0; i < k; i++) {
                    for (int boxVal : boxes[i]) {
                        Node currNode = nodes.get(boxVal);
                        long nextValLong = boxVal + delta[i];
                        if ((int) nextValLong != nextValLong) {
                            continue;
                        }
                        int nextVal = (int) nextValLong;
                        if (nodes.containsKey(nextVal)) {
                            Node nextNode = nodes.get(nextVal);
                            if (nextNode.boxId != i || nextVal == boxVal) {
                                currNode.next = nextNode;
                                remaining.add(currNode);
                            }
                        }

                    }
                }
                /* Extract Cycles */
                int max = (1 << k) - 1;
                HashMap<Integer, ArrayList<Node>> cycles = new HashMap<>();
                while (remaining.size() > 0) {
                    /* Get element */
                    Node top = remaining.iterator().next();
                    HashSet<Node> visited = new HashSet<>();
                    /* DFS from this element */
                    Node curr = top;
                    while (true) {
                        /* Process curr */
                        if (!visited.contains(curr)) {
                            visited.add(curr);
                            remaining.remove(curr);
                        } else {
                            /* Back edge */
                            int cycleID = 0;
                            ArrayList<Node> cycle = new ArrayList<>();
                            /* Extract cycle */
                            Node loop = curr;
                            boolean validCycle = true;
                            do {
                                if ((cycleID >> loop.boxId & 1) == 1) {
                                    validCycle = false;
                                    /* Dont add cycle, these nodes are already visited,
                                    so they will be removed anyways*/
                                    break;
                                }
                                cycleID |= 1 << loop.boxId;
                                cycle.add(loop);
                                loop = loop.prev;
                            } while (loop != curr && loop != null);
                            if (validCycle) {
                                cycles.putIfAbsent(cycleID, cycle);
                            }
                            /* Exit DFS */
                            break;
                        }
                        /* Check & Move pointer */
                        if (curr.next != null) {
                            /* set prev for this run */
                            curr.next.prev = curr;
                            /* move pointer */
                            curr = curr.next;
                        } else {
                            break;
                        }
                    }
                }
                /* Calculate Set Cover */
                boolean[] covered = new boolean[max + 1];
                int[] subCover = new int[max + 1];
                covered[0] = true;
                for (int i = 0; i <= max; i++) {
                    for (int j = i; j > 0; j = (j - 1) & i) {
                        if (cycles.containsKey(j) && covered[remove(i, j)]) {
                            subCover[i] = j;
                            covered[i] = true;
                            break;
                        }
                    }
                }
                /* Print Solution */
                if (covered[max]) {
                    out.println("Yes");
                    Queue<Integer> bfs = new LinkedList<>();
                    ArrayList<Node> solution = new ArrayList<>();
                    bfs.add(max);
                    while (bfs.size() > 0) {
                        int top = bfs.poll();
                        if (cycles.containsKey(top)) {
                            solution.addAll(cycles.get(top));
                        } else {
                            int sc = subCover[top];
                            bfs.add(sc);
                            bfs.add(remove(top, sc));
                        }
                    }
                    int[] resultVal = new int[k];
                    int[] resultId = new int[k];
                    for (Node p : solution) {
                        resultVal[p.boxId] = p.value;
                        resultId[p.next.boxId] = p.boxId + 1;
                    }
                    for (int i = 0; i < k; i++) {
                        out.println(resultVal[i] + " " + resultId[i]);
                    }
                } else {
                    out.println("No");
                }
            }

        }


        private int remove(int i, int j) {
            return i & (~j);
        }


        private static class Node {
            int value;
            int boxId;
            Node next;
            Node prev;

            Node(int value, int boxId) {
                this.value = value;
                this.boxId = boxId;
            }

            @Override
            public boolean equals(Object o) {
                if (this == o) {
                    return true;
                }
                if (o == null) {
                    return false;
                }
                Node node = (Node) o;
                return value == node.value;
            }

            @Override
            public int hashCode() {
                return value;
            }

            @Override
            public String toString() {
                return (value % 1000) + "[" + boxId + "]";
            }
        }

    }

    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());
        }

        public long nextLong() {
            return Long.parseLong(next());
        }

        public String nextLine() {
            try {
                return reader.readLine();
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }

        public int[] nextIntArray(int n) {
            int[] arr = new int[n];
            for (int i = 0; i < n; i++) {
                arr[i] = nextInt();
            }
            return arr;
        }

        public long[] nextLongArray(int n) {
            long[] arr = new long[n];
            for (int i = 0; i < n; i++) {
                arr[i] = nextLong();
            }
            return arr;
        }

    }

}
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.