如何使用Java模拟退火算法优化Hash函数

public class FindBestDistribution {

    static final int NUM_OF_SYMBOLS = 5000;
    static final int NUM_OF_BUCKETS = 15;

    public static void main(String[] args) {
        // 生成样本
        IntStream ints = ThreadLocalRandom.current().ints(0, 1000);
        PrimitiveIterator.OfInt iterator = ints.iterator();
        
        Map<String,Integer> symbolAndCount = new TreeMap<>();
        for (int i=0; i<NUM_OF_SYMBOLS; i++) {
            symbolAndCount.put(Integer.toHexString(i).toUpperCase(), iterator.next());
        }

        // 按照 symbol 划分每个桶的数量
        TreeMap<String, Integer> distribution = findBestDistribution(symbolAndCount);

        // 测试效果
        int[] buckets = new int[NUM_OF_BUCKETS];
        for (Map.Entry<String, Integer> entry : symbolAndCount.entrySet()) {
            Map.Entry<String, Integer> floor = distribution.floorEntry(entry.geTKEy());
            int bucketIndex = floor == null ? 0 : floor.getValue();
            buckets[bucketIndex] += entry.getValue();
        }

        System.out.printf("buckets: %s\n", Arrays.toString(buckets));
    }

    public static TreeMap<String, Integer> findBestDistribution(Map<String,Integer> symbolAndCount) {

        // 每个桶均匀分布的情况（最优情况）
        int avg = symbolAndCount.values().stream().mapToInt(Integer::intValue).sum() / NUM_OF_BUCKETS;

        // 尝试将 symbol 放入不同的桶
        int bucketIdx = 0;
        int[] buckets = new int[NUM_OF_BUCKETS];
        String[] bulkheads = new String[NUM_OF_BUCKETS-1];
        for (Map.Entry<String, Integer> entry : symbolAndCount.entrySet()) {

            // 如果首个 symbol 数据量过大，则分配给其一个独立的桶
            int count = entry.getValue();
            if (count / 2 > avg && bucketIdx == 0 && buckets[0] == 0) {
                buckets[bucketIdx] += count;
                continue;
            }

            // 评估将 symbol 放入桶后的效果
            // 1. 如果桶中的数量更接近期望，则将其放入当前桶中
            // 2. 如果桶中的数量更远离期望，则将其放入下个桶中
            double before = Math.abs(buckets[bucketIdx] - avg);
            double after = Math.abs(buckets[bucketIdx] + count - avg);
            if (after > before && bucketIdx < buckets.length - 1) {
                bulkheads[bucketIdx++] = entry.getKey();
            }

            buckets[bucketIdx] += count;
        }

        System.out.printf("expectation: %d\n", avg);
        System.out.printf("bulkheads: %s\n", Arrays.toString(bulkheads));

        TreeMap<String,Integer> distribution = new TreeMap<>();
        for (int i=0; i<bulkheads.length; i++) {
            distribution.put(bulkheads[i], i+1);
        }
        return distribution;
    }
}

int BKDRHash(char[] value, int seed) {
    int hash = 0;
    for (int i = 0; i < value.length; i++) {
        hash = hash * seed + value[i];
    }
    return hash & 0x7fffffff;
}

ouble calculateDivergence(long[] bucket, long expectation) {
    long divergence = 0;
    for (int i=0; i<bucket.length; i++) {
        final long a = bucket[i];
        final long b = (a - expectation) * (a - expectation);
        for (int j=i+1; j<bucket.length; j++) {
            long c = (a - bucket[j]) * (a - bucket[j]);
            divergence += Math.max(b, c);
        }
    }
    return divergence; // the less the better
}

public abstract class SimulatedAnnealing<X> {

    protected final int numberOfIterations;    // stopping condition for simulations

    protected final double coolingRate;        // the percentage by which we reduce the temperature of the system
    protected final double initialTemperature; // the starting energy of the system
    protected final double minimumTemperature; // optional stopping condition

    protected final long simulationTime;       // optional stopping condition
    protected final int detectionInterval;     // optional stopping condition

    protected SimulatedAnnealing(int numberOfIterations, double coolingRate) {
        this(numberOfIterations, coolingRate, 10000000, 1, 0, 0);
    }

    protected SimulatedAnnealing(int numberOfIterations, double coolingRate, double initialTemperature, double minimumTemperature, long simulationTime, int detectionInterval) {
        this.numberOfIterations = numberOfIterations;
        this.coolingRate = coolingRate;
        this.initialTemperature = initialTemperature;
        this.minimumTemperature = minimumTemperature;
        this.simulationTime = simulationTime;
        this.detectionInterval = detectionInterval;
    }

    protected abstract double score(X currentSolution);

    protected abstract X neighbourSolution(X currentSolution);

    public X simulateAnnealing(X currentSolution) {

        final long startTime = System.currentTimeMillis();

        // Initialize searching
        X bestSolution = currentSolution;
        double bestScore = score(bestSolution);
        double currentScore = bestScore;

        double t = initialTemperature;
        for (int i = 0; i < numberOfIterations; i++) {
            if (currentScore < bestScore) {
                // If the new solution is better, accept it unconditionally
                bestScore = currentScore;
                bestSolution = currentSolution;
            } else {
                // If the new solution is worse, calculate an acceptance probability for the worse solution
                // At high temperatures, the system is more likely to accept the solutions that are worse
                boolean rejectWorse = Math.exp((bestScore - currentScore) / t) < Math.random();
                if (rejectWorse || currentScore == bestScore) {
                    currentSolution = neighbourSolution(currentSolution);
                    currentScore = score(currentSolution);
                }
            }

            // Stop searching when the temperature is too low
            if ((t *= coolingRate) < minimumTemperature) {
                break;
            }

            // Stop searching when simulation time runs out
            if (simulationTime > 0 && (i+1) % detectionInterval == 0) {
                if (System.currentTimeMillis() - startTime > simulationTime)
                    break;
            }
        }

        return bestSolution;
    }
}


@Data
public class SimulatedAnnealingHashing extends SimulatedAnnealing<HashingSolution> {

    private static final int DISTRIBUTION_BATCH = 100;
    static final int SEARCH_BATCH = 200;

    private final int[] hashCodes = new int[SEARCH_BATCH];
    private final long[][] buckets = new long[SEARCH_BATCH][];

    @Data
    public class HashingSolution {

        private final int begin, range; // the begin and range for searching
        private int bestSeed;     // the best seed found in this search
        private long bestScore;   // the score corresponding to bestSeed

        private long calculateDivergence(long[] bucket) {
            long divergence = 0;
            for (int i=0; i<bucket.length; i++) {
                final long a = bucket[i];
                final long b = (a - expectation) * (a - expectation);
                for (int j=i+1; j<bucket.length; j++) {
                    long c = (a - bucket[j]) * (a - bucket[j]);
                    divergence += Math.max(b, c);
                }
            }
            return divergence; // the less the better
        }

        private HashingSolution solve() {

            if (range != hashCodes.length) {
                throw new IllegalStateException();
            }

            for (int i=0; i<range; i++) {
                Arrays.fill(buckets[i], hashCodes[i] = 0);
            }

            for (KeyDistribution[] bucket : distributions) {
                for (KeyDistribution distribution : bucket) {
                    Hashing.BKDRHash(distribution.getKey(), begin, hashCodes);
                    for (int k = 0; k< hashCodes.length; k++) {
                        int n = hashCodes[k] % buckets[k].length;
                        buckets[k][n] += distribution.getCount();
                    }
                }
            }

            int best = -1;
            long bestScore = Integer.MAX_VALUE;
            for (int i = 0; i< buckets.length; i++) {
                long score = calculateDivergence(buckets[i]);
                if (i == 0 || score < bestScore) {
                    bestScore = score;
                    best = i;
                }
            }

            if (best < 0) {
                throw new IllegalStateException();
            }

            this.bestScore = bestScore;
            this.bestSeed = begin + best;
            return this;
        }

        @Override
        public String toString() {
            return String.fORMat("(seed:%d, score:%d)", bestSeed, bestScore);
        }
    }

    private final KeyDistribution[][] distributions; // key and its count（2-dimensional array for better performance）
    private final long expectation;  // the expectation count of each bucket
    private final int searchOutset;
    private int searchMin, searchMax;

    
    public SimulatedAnnealingHashing(Map<String, Integer> keyAndCounts, int numOfBuckets) {
        super(100000000, .9999);
        distributions = buildDistribution(keyAndCounts);
        long sum = 0;
        for (KeyDistribution[] batch : distributions) {
            for (KeyDistribution distribution : batch) {
                sum += distribution.getCount();
            }
        }
        this.expectation = sum / numOfBuckets;
        this.searchOutset = 0;
        for (int i = 0; i< buckets.length; i++) {
            buckets[i] = new long[numOfBuckets];
        }
    }

    
    private SimulatedAnnealingHashing(SimulatedAnnealingHashing prototype, int searchOutset, long simulationTime) {
        super(prototype.numberOfIterations, prototype.coolingRate, prototype.initialTemperature, prototype.minimumTemperature,
                simulationTime, 10000);
        distributions = prototype.distributions;
        expectation = prototype.expectation;
        for (int i = 0; i< buckets.length; i++) {
            buckets[i] = new long[prototype.buckets[i].length];
        }
        this.searchOutset = searchOutset;
        this.searchMax = searchMin = searchOutset;
    }

    @Override
    public String toString() {
        return String.format("expectation: %d, outset:%d, search(min:%d, max:%d)", expectation, searchOutset, searchMin, searchMax);
    }

    private KeyDistribution[][] buildDistribution(Map<String, Integer> symbolCounts) {
        int bucketNum = symbolCounts.size() / DISTRIBUTION_BATCH + Integer.signum(symbolCounts.size() % DISTRIBUTION_BATCH);
        KeyDistribution[][] distributions = new KeyDistribution[bucketNum][];

        int bucketIndex = 0;
        List<KeyDistribution> batch = new ArrayList<>(DISTRIBUTION_BATCH);
        for (Map.Entry<String, Integer> entry : symbolCounts.entrySet()) {
            batch.add(new KeyDistribution(entry.getKey().toCharArray(), entry.getValue()));
            if (batch.size() == DISTRIBUTION_BATCH) {
                distributions[bucketIndex++] = batch.toArray(new KeyDistribution[0]);
                batch.clear();
            }
        }
        if (batch.size() > 0) {
            distributions[bucketIndex] = batch.toArray(new KeyDistribution[0]);
            batch.clear();
        }
        return distributions;
    }

    @Override
    protected double score(HashingSolution currentSolution) {
        return currentSolution.solve().bestScore;
    }

    @Override
    protected HashingSolution neighbourSolution(HashingSolution currentSolution) {
        // The default range of neighbourhood is [-100, 100]
        int rand = ThreadLocalRandom.current().nextInt(-100, 101);
        int next = currentSolution.begin + rand;
        searchMin = Math.min(next, searchMin);
        searchMax = Math.max(next, searchMax);
        return new HashingSolution(next, currentSolution.range);
    }

    public HashingSolution solve() {
        searchMin = searchMax = searchOutset;
        HashingSolution initialSolution = new HashingSolution(searchOutset, SEARCH_BATCH);
        return simulateAnnealing(initialSolution);
    }

    public SimulatedAnnealingHashing derive(int searchOutset, long simulationTime) {
        return new SimulatedAnnealingHashing(this, searchOutset, simulationTime);
    }
}

@Data
@Slf4j
public class HashingSeedCalculator {

    
    private class HashingSeedCalculatorSearchTask extends RecursiveTask<HashingSolution> {

        private SimulatedAnnealingHashing simulation;
        private final int level;
        private final int center, range;

        private HashingSeedCalculatorSearchTask() {
            this.center = 0;
            this.range = Integer.MAX_VALUE / SimulatedAnnealingHashing.SEARCH_BATCH;
            this.level = traversalDepth;
            this.simulation = hashingSimulation;
        }

        private HashingSeedCalculatorSearchTask(HashingSeedCalculatorSearchTask parent, int center, int range) {
            this.center = center;
            this.range = range;
            this.level = parent.level - 1;
            this.simulation = parent.simulation;
        }

        @Override
        protected HashingSolution compute() {
            if (level == 0) {
                long actualCenter = center * SimulatedAnnealingHashing.SEARCH_BATCH;
                log.info("Searching around center {}", actualCenter);
                HashingSolution solution = simulation.derive(center, perShardRunningMills).solve();
                log.info("Searching around center {} found {}", actualCenter, solution);
                return solution;
            } else {
                int halfRange = range / 2;
                int leftCenter = center - halfRange, rightCenter = center + halfRange;
                ForkJoinTask<HashingSolution> leftTask = new HashingSeedCalculatorSearchTask(this, leftCenter, halfRange).fork();
                ForkJoinTask<HashingSolution> rightTask = new HashingSeedCalculatorSearchTask(this, rightCenter, halfRange).fork();
                HashingSolution left = leftTask.join();
                HashingSolution right = rightTask.join();
                return left.getBestScore() < right.getBestScore() ? left : right;
            }
        }
    }

    private final int poolParallelism;
    private final int traversalDepth;
    private final long perShardRunningMills;
    private final SimulatedAnnealingHashing hashingSimulation;

    
    public HashingSeedCalculator(int numberOfShards, int totalRunningHours, Map<String, Integer> symbolCounts, int numOfBuckets) {
        int n = (int) (Math.log(numberOfShards) / Math.log(2));
        if (Math.pow(2, n) != numberOfShards) {
            throw new IllegalArgumentException();
        }
        this.traversalDepth = n;
        this.poolParallelism = Math.max(ForkJoinPool.getCommonPoolParallelism() / 3 * 2, 1); // conservative estimation for parallelism
        this.perShardRunningMills = TimeUnit.HOURS.toMillis(totalRunningHours * poolParallelism) / numberOfShards;
        this.hashingSimulation = new SimulatedAnnealingHashing(symbolCounts, numOfBuckets);
    }

    @Override
    public String toString() {
        int numberOfShards = (int) Math.pow(2, traversalDepth);
        int totalRunningHours = (int) TimeUnit.MILLISECONDS.toHours(perShardRunningMills * numberOfShards) / poolParallelism;
        return "HashingSeedCalculator(" +
                "numberOfShards: " + numberOfShards +
                ", perShardRunningMinutes: " + TimeUnit.MILLISECONDS.toMinutes(perShardRunningMills) +
                ", totalRunningHours: " + totalRunningHours +
                ", poolParallelism: " + poolParallelism +
                ", traversalDepth: " + traversalDepth + ")";
    }

    public synchronized HashingSolution searchBestSeed() {
        long now = System.currentTimeMillis();
        log.info("SearchBestSeed start");
        ForkJoinTask<HashingSolution> root = new HashingSeedCalculatorSearchTask().fork();
        HashingSolution initSolution = hashingSimulation.derive(0, perShardRunningMills).solve();
        HashingSolution bestSolution = root.join();
        log.info("Found init solution {}", initSolution);
        log.info("Found best solution {}", bestSolution);
        if (initSolution.getBestScore() < bestSolution.getBestScore()) {
            bestSolution = initSolution;
        }
        long cost = System.currentTimeMillis() - now;
        log.info("SearchBestSeed finish (cost:{}ms)", cost);
        return bestSolution;
    }

}