Javaの優先順位付きキュー(PriorityQueue)の使い方について(AtCoderRegulerContest 028 B問題)

JavaでPriorityQueueを初めて使ったのですが、なかなか理解できなかったんで、知識の整理も兼ねて書きます。
ネット上で探しても、初心者の僕に対して、これでばっちり！という記事がなかったので。

ちなみに、使うきっかけとなった問題は練習として解くことにしたARC#028のB問題です。
満点解法が分からなかったので解説スライドを見てみると、このPriorityQueueが出てきました。

概念としてのだいたいの説明をwikiで読んで、あとはこのサイトを見ながら、色々試してみる。ドキドキ。

まずは、一番シンプルなパターン？から。
引数なしでキューを生成してみます。そうすると『自然順序付け に従って要素を順序付けする、デフォルトの初期容量 (11) を持つ PriorityQueue を作成します。』とのこと。自然順序付けって、単純に数の大小ってこと？
やってみました。

import java.util.*;
public class Main {
	public static void main(String[] args){

 		Queue queue = new PriorityQueue();
        queue.add(3);
        System.out.println(queue.peek());
        queue.add(2);
        System.out.println(queue.peek());
        queue.add(4);
        System.out.println(queue.peek());
        queue.add(1);
        System.out.println(queue.peek());
	}
}

addで引数をキューに追加、peekで『キューの先頭を返す(削除はしない)』という命令です。
これを、いつもお世話になっております、ideoneで実行してみる。以降のものも、ideoneで実行しています。
結果はこんな感じ。

Success	time: 0.07 memory: 380224 signal:0
3
2
2
1

ということは、自然順序付けでは、小さい数ほど先頭に行くってことですね。なるほど。

実数でも出来るのか？という実験。

import java.util.*;
public class Main {
	public static void main(String[] args){

 		Queue queue = new PriorityQueue();
        queue.add(3.3);
        System.out.println(queue.peek());
        queue.add(3.2);
        System.out.println(queue.peek());
        queue.add(3.4);
        System.out.println(queue.peek());
        queue.add(3.1);
        System.out.println(queue.peek());
	}
}

結果。

Success	time: 0.08 memory: 380224 signal:0
3.3
3.2
3.2
3.1

できる。

整数と実数をまぜると？

import java.util.*;
public class Main {
	public static void main(String[] args){

 		Queue queue = new PriorityQueue();
        queue.add(3.3);
        System.out.println(queue.peek());
        queue.add(3);
        System.out.println(queue.peek());
        queue.add(3.4);
        System.out.println(queue.peek());
        queue.add(2.9);
        System.out.println(queue.peek());
	}
}

結果。

Runtime error	time: 0.07 memory: 380224 signal:-1
3.3

ここで止まりました。
ということは、同じ型じゃないとダメってことですかね。へぇ。

次に、今回の問題ではキュー内で一番大きい数を返して欲しいので、それが可能か実験してみます。
比較関数を実装して、うんちゃらかんちゃら、というのが出来るようです。
このサイトを参考にさせて頂きました。

とりあえず、上のサイトの、MyComparatorを整数に関するものに変更して、実行してみます。

import java.util.*;

class MyComparator implements Comparator {
 
    public int compare(Object obj1, Object obj2) {
         
        int num1 = (int)obj1;
        int num2 = (int)obj2;
         
        if(num1 > num2) {
            return 1;
        } else if(num1 < num2) {
            return -1;
        } else{
            return 0;
        }
    }
}

public class Main {
	public static void main(String[] args){
 		Queue queue = new PriorityQueue(4,new MyComparator());
        queue.add(3);
        System.out.println(queue.peek());
        queue.add(2);
        System.out.println(queue.peek());
        queue.add(4);
        System.out.println(queue.peek());
        queue.add(1);
        System.out.println(queue.peek());
	}
}

結果はこうなりました。

Success	time: 0.08 memory: 380160 signal:0
3
2
2
1

つまり、上で書いたMyComparatorが、"自然順序"であるということですね。
ちなみに、実行するコードを書くまで、10分くらいかかりました。publicとかclassとかnewとか、ややこしいです。しばくぞ

さて、では、先頭が一番大きい数になるように書き直してみます。とはいっても、大小関係を逆にするだけですね。

import java.util.*;

class MyComparator implements Comparator {
 
    public int compare(Object obj1, Object obj2) {
         
        int num1 = (int)obj1;
        int num2 = (int)obj2;
         
        if(num1 < num2) {
            return 1;
        } else if(num1 > num2) {
            return -1;
        } else{
            return 0;
        }
    }
}

public class Main {
	public static void main(String[] args){
 		Queue queue = new PriorityQueue(4,new MyComparator());
        queue.add(3);
        System.out.println(queue.peek());
        queue.add(2);
        System.out.println(queue.peek());
        queue.add(4);
        System.out.println(queue.peek());
        queue.add(1);
        System.out.println(queue.peek());
	}
}

実行！

Success	time: 0.07 memory: 380224 signal:0
3
3
4
4

おお、できてる。少し感動です。
ここまでくれば、今回の問題もすぐ解けそうですね。

というわけで、その後頑張って、なんとか解けました。こんなコードになりました。

import java.util.*;

class MyComparator implements Comparator {
 
    public int compare(Object obj1, Object obj2) {
         
        int num1 = (int)obj1;
        int num2 = (int)obj2;
         
        if(num1 < num2) {
            return 1;
        } else if(num1 > num2) {
            return -1;
        } else{
            return 0;
        }
    }
}

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

		int nowKth;

 		int[] persons = new int[n];

 		// ranking[i]には、全体で(i+1)番目に若い人が、何位だったかを保存していく。
 		int[] ranking = new int[n];

 		for (int i = 0; i < n; i++) {
 			persons[i] = sc.nextInt();
 		}

 		Queue queue = new PriorityQueue(k,new MyComparator());

 		for (int i = 0; i < k; i++) {
 			ranking[persons[i]-1] = i + 1;
 			queue.add(persons[i]);
 		}

 		nowKth = (int)queue.peek();
 		System.out.println(ranking[nowKth - 1]);

 		for (int i = k; i < n; i++) {
 			if (nowKth > persons[i]) {
 				ranking[persons[i]-1] = i + 1;
 				queue.add(persons[i]);
 				queue.poll();
 				nowKth = (int)queue.peek();
 			}
  			System.out.println(ranking[nowKth - 1]);
 		}
	}
}

ちなみに、提出したらACになったものの、実行時間は1889ms。大分ギリギリですね。
他の早いひとのコードもちょっと見たのですが、あんまり読めなかったです。えへへ

解いている途中で気づいた注意点としては、キューからpeek()で引っ張ってきたものは、Object型？というものになっているということです。
例えば、以下のコードを実行すると

import java.util.*;
public class Main {
	public static void main(String[] args){
 		Queue queue = new PriorityQueue();
        queue.add(3);
        System.out.println(queue.peek() - 1);
	}
}

こうなります。

Compilation error	time: 0.1 memory: 380672 signal:0
Main.java:6: error: bad operand types for binary operator '-'
        System.out.println(queue.peek() - 1);
                                        ^
  first type:  Object
  second type: int
Note: Main.java uses unchecked or unsafe operations.
Note: Recompile with -Xlint:unchecked for details.
1 error

このエラーを防ぐには、単に『(int)queue.peek()』という風にキャストしてやればOKです。
もう、ややこしいなこういうの。

はい、今回は以上です。
初めてのデータ構造でした。
これからも頑張るぞー