Bear and Displayed Friends

Introduction

Limak is a little polar bear. He loves connecting with other bears via social networks. He has n friends and his relation with the i-th of them is described by a unique integer ti. The bigger this value is, the better the friendship is. No two friends have the same value ti.

Spring is starting and the Winter sleep is over for bears. Limak has just woken up and logged in. All his friends still sleep and thus none of them is online. Some (maybe all) of them will appear online in the next hours, one at a time.

The system displays friends who are online. On the screen there is space to display at most k friends. If there are more than k friends online then the system displays only k best of them — those with biggest ti.

Your task is to handle queries of two types:

"1 id" — Friend id becomes online. It's guaranteed that he wasn't online before. "2 id" — Check whether friend id is displayed by the system. Print "YES" or "NO" in a separate line. Are you able to help Limak and answer all queries of the second type?

Solutions

Idea

Friends are associated with a specific friendship score. As they come online, you want to keep track of the k highest scored friends. An easy way to do that is to use a TreeMap data structure, with (score, friend ID) as the key-value pair. It is safe to use the score as a key because you are guaranteed that scores are unique among all friends. As friends come online, insert them into the TreeMap. If the TreeMap size grows larger than k, remove the lowest-scoring friend.

Runtime

Insertion into a TreeMap is O(logn), so query 1 for friends coming online would be O(logn). Searching for an item in a TreeMap is also O(logn), so query 2 which queries whether a friend is online would also be O(logn).

Code

Solution - Java

import java.util.Scanner;
import java.util.TreeMap;

/**
 * Created by alex on 3/28/16.
 * Problem found at http://codeforces.com/contest/658/problem/B
 */
public class BearAndDisplayedFriends {

    public static TreeMap<Integer, Integer> onlineFriends = new TreeMap<>();

    public static void main(String args[]){
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        int k = in.nextInt();
        int q = in.nextInt();

        int[] friends = new int[n];
        for(int i=0; i<n; i++){
            friends[i] = in.nextInt();
        }

        for(int i=0; i<q; i++){
            int type = in.nextInt();
            int id = in.nextInt();
            if(type == 1){
                onlineFriends.put(friends[id-1], id);
                if(onlineFriends.size() > k){
                    onlineFriends.pollFirstEntry().getKey();
                }
            }
            else{
                solve(id);
            }
        }
    }

    public static void solve(int id){
        if(onlineFriends.containsValue(id)){
            System.out.println("YES");
        }
        else{
            System.out.println("NO");
        }
    }
}