leetcode-go/basis/algorithms/skiplist/java/SkipList.java

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.Stack;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

public class SkipList<T extends Comparable<T>> {
    private final ThreadLocalRandom random;

    private final int maxLevel;
    private final AtomicInteger curLevel;
    private final double p;
    private final ListNode<T> head;

    public SkipList(int maxLevel, double p) {
        this.maxLevel = maxLevel;
        this.p = p;
        this.random = ThreadLocalRandom.current();
        this.curLevel = new AtomicInteger(0);
        this.head = new ListNode<T>(null, maxLevel);
    }

    public void add(T key) {
        insert(key);
    }

    public void insert(T key) {
        Stack<ListNode<T>> update = new Stack<>();
        ListNode<T> curr = head;
        for (int level = curLevel.get(); level >= 0; level--) {
            while (curr.canMoveForward(key, level)) {
                curr = curr.forward.get(level).get();
            }
            update.add(curr);
        }
        ListNode<T> next = curr.forward.get(0).get();
        if (next != null && key.compareTo(next.key) == 0) {
            // Update value
            return;
        }
        ListNode<T> down = null;
        int insertLevel = randomLevel();
        for (int level = 0; level <= insertLevel; level++) {
            curr = update.isEmpty() ? head : update.pop();
            ListNode<T> up = insertAfter(key, curr, level);
            if (level > 0) {
                up.forward.get(level - 1).set(down);
            }
            down = up;
        }
        while (true) {
            int expectedLevel = curLevel.get();
            if (expectedLevel >= insertLevel || curLevel.compareAndSet(expectedLevel, insertLevel)) {
                break;
            }
        }
    }

    public boolean contains(T key) {
        ListNode<T> curr = head;
        for (int level = curLevel.get(); level >= 0; level--) {
            while (curr.canMoveForward(key, level)) {
                curr = curr.forward.get(level).get();
            }
        }
        curr = curr.forward.get(0).get();
        return curr != null && key.compareTo(curr.key) == 0;
    }

    @Override
    public String toString() {
        StringBuilder sb = new StringBuilder();
        sb.append("current level: ").append(curLevel.get()).append("\n");
        for (int i = curLevel.get(); i >= 0; i--) {
            ListNode<T> curr = head.forward.get(i).get();
            while (curr != null) {
                sb.append(curr.key.toString()).append(" ");
                curr = curr.forward.get(i).get();
            }
            sb.append("\n");
        }
        return sb.toString();
    }

    private ListNode<T> insertAfter(T key, ListNode<T> before, int level) {
        ListNode<T> after = new ListNode<>(key, level);
        while (true) {
            while (before.canMoveForward(key, level)) {
                before = before.forward.get(level).get();
            }
            ListNode<T> next = before.forward.get(level).get();
            after.forward.get(level).set(next);
            if (before.forward.get(level).compareAndSet(next, after)) {
                break;
            }
        }
        return after;
    }

    private int randomLevel() {
        int level = 0;
        while (level < maxLevel) {
            if (random.nextDouble() < p) {
                level++;
            } else {
                break;
            }
        }
        return level;
    }

    public static void main(String[] args) throws InterruptedException {
        SkipList<Integer> skipList = new SkipList<>(20, 0.2);
//        Set<Integer> skipList = new ConcurrentSkipListSet<>();
        int threads = 1000;
        int loop = 100;
        ExecutorService executor = Executors.newCachedThreadPool();
        CountDownLatch cd = new CountDownLatch(threads);
        for (int i = 0; i < threads; i++) {
            final int initJ = i;
            executor.submit(() -> {
                for (int j = initJ; j < threads * loop; j += threads) {
                    skipList.add(j);
                }
                cd.countDown();
            });
        }
        cd.await();
        executor.shutdown();
        for (int i = 0; i < threads * loop; i++) {
            if (!skipList.contains(i)) {
                System.out.println("Concurrent error: " + i);
            }
        }
    }

    public static class ListNode<T extends Comparable<T>> {
        private List<AtomicReference<ListNode<T>>> forward;
        private T key;

        public ListNode(T key, int level) {
            this.key = key;
            this.forward = new ArrayList<>(level + 1);
            for (int i = 0; i <= level; i++) {
                this.forward.add(new AtomicReference<>());
            }
        }

        public boolean canMoveForward(T key, int level) {
            return this.forward.get(level).get() != null &&
                    key.compareTo(this.forward.get(level).get().key) > 0;
        }

        @Override
        public String toString() {
            return "ListNode{" +
                    "forward=" + forward +
                    ", key=" + key +
                    '}';
        }
    }
}