優先度付きキュー(Priority Queue)の実装 in Java

プログラミング能力向上のため、競技プログラミングの問題を多数収録しているAizu Online Judgeの問題を少しずつ解いています。

以前解けなかった優先度付きキューの問題を解きなおしました。二分ヒープで実装されているので、最大値の取得はO(1)、要素の追加を最悪計算量O(log n)で実行できます。
優先度付きキュー | アルゴリズムとデータ構造 | Aizu Online Judge
(参考:二分ヒープ https://en.wikipedia.org/wiki/Binary_heap)

リファクタリングで改善したのは、リストを利用していたところを、配列を利用することにしただけですが…。

企業での実際の現場ではこんな典型的なデータ構造は自分でコーディングせず、java.util.PriorityQueueのようなライブラリを使うのがベストですね！^^;

ProgrammingContests/ALDS1_9_C.java at master · toricor/ProgrammingContests · GitHub

import java.io.IOException;
import java.io.BufferedReader;
import java.io.InputStreamReader;
//AC
public class Main {
    private static int[] A = new int[2000000];
    private static int heapSize = 0;
    private static void maxHeapify(int k){
        int l = 2 * k;
        int r = 2 * k + 1;
        int largest = 0;
      
        if ((l <= (heapSize)) && (A[l] > A[k])){
            largest = l;
        }else{
            largest = k;
        }
        if ((r <= (heapSize)) && (A[r] > A[largest])){
            largest = r;
        }
        if (largest != k){
            swap(k, largest);
            maxHeapify(largest);
        }
    }
    private static void swap(int i, int j) {
        int tmp = A[i];
        A[i] = A[j];
        A[j] = tmp;
    }
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        while (true){
            String[] line = br.readLine().split(" ");
            String order = line[0];
            if (order.equals("insert")){
                heapSize++;
                int num = Integer.parseInt(line[1]);
                A[heapSize] = num;
                int i = heapSize;
                while (i > 1 && A[i/2] < A[i]) {
                    swap(i, i/2);
                    i = i/2;
                }
            }else if (order.equals("extract")){
                int max = A[1];
                A[1] = 0;
                swap(1, heapSize);
                heapSize--;
                maxHeapify(1);
                System.out.println(max);
            }else{ // end
                break;
            }
        }
    }
}

以下はArrayListを使用してPriority Queueを実装した例です。テストケースのうち、最も大きい最後のケースを制限時間内に通過できませんでした。

import java.io.IOException;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.List;
import java.util.ArrayList;
//Listではラストケースを通過できなかった
public class Main {
    static List<Integer> A = new ArrayList<Integer>();
    static void maxHeapify(List<Integer> A, int k){
        int l = 2 * k;
        int r = 2 * k + 1;
        int largest = 0;
    
        if ((l <= (A.size() - 1)) && (A.get(l) > A.get(k))){
            largest = l;
        }else{
            largest = k;
        }
        if ((r <= (A.size() - 1)) && (A.get(r) > A.get(largest))){
            largest = r;
        }
        if (largest != k){
            int tmp = A.get(k);
            A.set(k, A.get(largest));
            A.set(largest, tmp);
            maxHeapify(A, largest);
        }
    }
    
    static void buildMaxHeap(List<Integer> A){
        for (int i=(A.size() - 1) / 2; i >= 0; i--){
            maxHeapify(A, i);
        }
    }        

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        
        while (true){
            String[] line = br.readLine().split(" ");
            String order = line[0];
            if (order.equals("insert")){
                int num = Integer.parseInt(line[1]);
                A.add(num);
            }else if (order.equals("extract")){
                buildMaxHeap(A);
                System.out.println(A.get(0));
                A.remove(0);
            }else{ // end
                break;
            }
        }
    }
}