Nearby

Nearby.java

import java.io.InputStream;
import java.io.OutputStream;
import java.io.PrintWriter;
import java.util.*;

/**
 * A class to implement http://www.quora.com/challenges#nearby
 * @author Kedar Mhaswade
 */
public class Nearby {

    static final Map<Integer, Topic> topics = new HashMap<>();//topic.id -> topic
    static int ntopics, nquestions, nqueries;

    public static void main(String args[]) {
        run(System.in, System.out);
    }

    static void run(InputStream is, OutputStream os) {
        Scanner s = new Scanner(is);
        Scanner ls = new Scanner(s.nextLine()).useDelimiter("\\s+");
        ntopics = ls.nextInt();
        nquestions = ls.nextInt();
        nqueries = ls.nextInt();
        //build topics first
        for (int i = 0; i < ntopics; i++) {
            ls = new Scanner(s.nextLine()).useDelimiter("\\s+");
            int id = ls.nextInt();
            double x = ls.nextDouble(), y = ls.nextDouble();
            topics.put(id, new Topic(id, x, y));
        }
        //add questions now
        for (int i = 0; i < nquestions; i++) {
            ls = new Scanner(s.nextLine()).useDelimiter("\\s+");
            int qid = ls.nextInt();
            int ts = ls.nextInt();
            for (int j = 0; j < ts; j++) {
                topics.get(ls.nextInt()).add(qid);
            }
        }
        //answer queries
        PrintWriter out = new PrintWriter(os);
        for (int i = 0; i < nqueries; i++) {
            String line = s.nextLine();
            ls = new Scanner(line).useDelimiter("\\s+");
            char c = ls.next().charAt(0);
            int nr = ls.nextInt(); // number of results;
            double x = ls.nextDouble(), y = ls.nextDouble();
            if (c == 't') {
                printTopics(x, y, nr, out);
            } else { // has to be 'q' for questions?
                printQuestions(x, y, nr, out);
            }
        }
        out.flush();
    }
    static void printTopics(double x, double y, int n, PrintWriter out) {
        PriorityQueue<Topic> minHeap = buildHeap(x, y);
        for (int i = 0; i < n && !minHeap.isEmpty(); i++) {
            Topic t = minHeap.remove();
            out.print(t.id);
            out.print(" ");
        }
        out.println();
    }
    static void printQuestions(double x, double y, int n, PrintWriter out) {
        PriorityQueue<Topic> minHeap = buildHeap(x, y);
        Set<Integer> collector = new LinkedHashSet<>();
        while (n > 0 && !minHeap.isEmpty()) {
            Topic t = minHeap.remove();
            PriorityQueue<Integer> qids = t.questions;
            while (!qids.isEmpty()) {
                int candidate = qids.remove();
                if (!collector.contains(candidate)) {
                    collector.add(candidate);
                    n -= 1;
                }
                if (n == 0)
                    break;
            }
        }
        for (int qid : collector) {
            out.print(qid);
            out.print(" ");
        }
        out.println();
    }
    static PriorityQueue<Topic> buildHeap(double x, double y) {
        Topic tt = new Topic(-1, x, y); //the test topic has a negative id, because only location is of essence
        PriorityQueue<Topic> minHeap = new PriorityQueue<>(100, new TopicDistanceIdComparator(tt));
        for (Map.Entry<Integer, Topic> entry : topics.entrySet()) {
            minHeap.add(entry.getValue());
        }
        return minHeap;
    }
    static class Topic {
        final int id;
        final double x, y;
        final PriorityQueue<Integer> questions = new PriorityQueue<>(100, ReverseIntComparator.INSTANCE);
        Topic(int id, double x, double y) {
            this.id = id;
            this.x = x;
            this.y = y;
        }
        void add(int qid) {
            questions.add(qid);
        }
        double from(Topic that) {
            double dx = this.x - that.x;
            double dy = this.y - that.y;
            return (dx*dx + dy*dy);
        }
    }

    static class TopicDistanceIdComparator implements Comparator<Topic> {
        final Topic test;
        TopicDistanceIdComparator(Topic test) {
            this.test = test;
        }
        @Override
        public int compare(Topic t1, Topic t2) {
            double d1 = t1.from(test);
            double d2 = t2.from(test);
            if (d1 == d2) {
                //order according to reverse natural order of topic ID's
                if (t1.id == t2.id)
                    return 0;
                return t1.id > t2.id ? -1 : 1;
            }
            return d1 < d2 ? -1 : 1;
        }
    }

    static class ReverseIntComparator implements Comparator<Integer> {
        private ReverseIntComparator() {}; //private
        public static ReverseIntComparator INSTANCE = new ReverseIntComparator();
        @Override
        public int compare(Integer o1, Integer o2) {
            if (o1 == o2)
                return 0;
            if (o1 > o2)
                return -1;
            return 1;
        }
    }
}