WaveletMatrix - nise-nabe/topcoder GitHub Wiki

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class WaveletTree {
    int[] a;
    boolean[][] b;

    public int access(int i) {
        return a[i];
    }

    /** 0 から i (exclusive) までに出現する x の数を返す */
    public int rank(int i, int x) {
        return rank(0, a.length, i, x, 0);
    }

    /** from から to をベースとして i までの x の depth 番目のビット数を返す*/
    private int rank(int from, int to, int i, int x, int depth) {
        if (depth >= b.length) {
            return i - from;
        }
        // row は逆に設置してあるので n として算出する
        int n = b.length - depth - 1;
        // 対象となる x の n depth 目の bit を出す
        boolean bb =  (x & 1 << n) >> n == 1;
        // from から to までの間にある 0 の数を出して 次の from または to の数を出す
        int med = from + rank(from, to, false, depth);
        // 次の i を計算する（範囲内における 1 の開始位置）
        i = rank(i, bb, depth) - rank(from, bb, depth);
        if (bb) { // 右側
            return rank(med, to, med + i, x, depth + 1);
        } else { // 左側
            // bb が 0 の場合は次のランクは左側で計算するので
            // from, from + (rank(to, 0, depth) - rank(from, 0, depth))
            return rank(from, from + med, from + i, x, depth + 1);
        }

    }

    /** from から to までの bb ビットの数を返す */
    private int rank(int from, int to, boolean bb, int depth) {
        return rank(to, bb, depth) - rank(from, bb, depth);

    }

    /** 0 から i までの bb ビットの数を返す */
    private int rank(int i, boolean bb, int depth) {
        int count = 0;
        for(int j = 0; j < i; ++j) {
            if (b[depth][j] == bb) {
                ++count;
            }
        }
        return count;
    }

    /** (i+1) 番目の c の出現位置を返す */
    public int select(int i, int x) {
        return 0;
    }

    public static WaveletTreeBuilder builder() {
        return new WaveletTreeBuilder();
    }
}

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class WaveletTreeBuilder {
    private int max, a[];

    public WaveletTreeBuilder max(int max) {
        this.max = max;
        return this;
    }
    public WaveletTreeBuilder a(int[] a) {
        this.a = a;
        return this;
    }

    public WaveletTree build() {
        WaveletTree wt = new WaveletTree();
        wt.a = a;
        // 高さ log2(max) * 数列の長さ
        int n = 32 - Integer.numberOfLeadingZeros(max);
        wt.b = new boolean[n][a.length];

        int[] aa = Arrays.copyOf(a, a.length);
        for (int i = n; i --> 0;) {
            List<Integer> zero = new ArrayList<>();
            List<Integer> one = new ArrayList<>();
            for (int j = 0; j < aa.length; ++j) {
                wt.b[n - i - 1][j] = (aa[j] & 1 << i) >> i == 1;
                // 0 の場合は左用、 1 の場合は右用に a [j] の値を振り分ける必要がある
                if ((aa[j] & 1 << i) >> i == 0) {
                    zero.add(aa[j]);
                } else {
                    one.add(aa[j]);
                }
            }
            List<Integer> list = new ArrayList<>();
            list.addAll(zero);
            list.addAll(one);
            aa = new int[aa.length];
            for (int k = 0; k < aa.length; ++k) {
                aa[k] = list.get(k);
            }
        }

        return wt;
    }

}

    public static void main(String[] args) {
        WaveletTree wt = WaveletTree.builder()
                                    .max(7)
                                    .a(new int[]{0,7,2,1,4,3,6,7,2,5,0,4,7,2,6,3})
                                    .build();

        for(int i = 0; i <= 16; ++i) {
            System.out.println(wt.rank(i, 7));
        }

    }