java不放回抽取_不等概率不放回

import java.util.ArrayList;

import java.util.List;

import java.util.Random;

/**

* 不等概率不放回的抽样类 使用方法：传入你的概率rates，以及需要抽取的样本个数k。假如传入的概率是：[1,2,3,4,5]，

* k为2，如果最后选择到的概率是1,3；那么返回的index为0(概率1的index),2(概率3的index)

*

* @author xutaoyang

*

*/

public class UnequalWithoutReplacementKRandom {

static Random rand = new Random();

/**

* 对外接口方法

*

* @param rates

*            概率

* @param k

*            目标样本的个数

* @return 命中的样本的在概率list中的index

*/

public static List randKWithoutReplacement(List rates, int k) {

if (null == rates || rates.isEmpty()) {

throw new RuntimeException("<> the rates list is null or empty");

}

if (k >= rates.size()) {

throw new RuntimeException("<> k is bigger than rates' size");

}

List nodes = new ArrayList(rates.size());

for (int index = 0; index < rates.size(); index++) {

nodes.add(new Node(rates.get(index), index));

}

List result = new ArrayList(k);

List heap = buildHeap(nodes);

for (int index = 0; index < k; index++) {

result.add(heapPop(heap));

}

return result;

}

private static List buildHeap(List nodes) {

List heap = new ArrayList(nodes.size() + 1);

heap.add(null);

for (int index = 0; index < nodes.size(); index++) {

heap.add(nodes.get(index));

}

for (int index = heap.size() - 1; index > 1; index--) {

double curTW = heap.get(index >> 1).totalWeight;

heap.get(index >> 1).totalWeight = curTW + heap.get(index).totalWeight;

}

return heap;

}

/** 关于double的计算都用+-x/了，那点误差就让它去吧，性能高很多啊 */

private static int heapPop(List heap) {

double gas = heap.get(1).totalWeight * rand.nextDouble();

int i = 1;

while (gas > heap.get(i).weight) {

gas = gas - heap.get(i).weight;

i <<= 1;

if (gas > heap.get(i).totalWeight) {

gas = gas - heap.get(i).totalWeight;

i++;

}

}

double weight = heap.get(i).weight;

int value = heap.get(i).value;

heap.get(i).weight = 0;

while (i > 0) {

heap.get(i).totalWeight = heap.get(i).totalWeight - weight;

i >>= 1;

}

return value;

}

private static class Node {

double weight;

int value;

double totalWeight;

public Node(double weight, int value) {

this.weight = weight;

this.value = value;

this.totalWeight = weight;

}

}

}