Implementing HashSet in Java

Implementing HashSet in Java.

 

public class HashSet<K> {
	public static class Node<K> {
		final K key;
		Node next;
		public Node(K key, Node next) {
			this.key = key;
			this.next = next;
		}
	}
	
	private int size = 0;
	private int capacity = 16;
	private Node[] tab;
	
	public HashSet(int capacity) {
		this.capacity = capacity; // should check range, make it 2^n
		tab = new Node[this.capacity];
	}
	
	private int indexOf(K key) {
		int hash = key == null ? 0 : key.hashCode();
		return hash % this.capacity;
		//return hash & (tab.length-1);
	}
	
	public boolean add(K key) {
		int idx = indexOf(key);
		Node<K> node = tab[idx];
		if(node == null) {
			tab[idx] = new Node(key, null);
		} else if(node.key == key || (key != null && key.equals(node.key))) {
			return false;
		} else {
			tab[idx] = new Node(key, node);
		}
		size++;
		return true;
	}
	
	public boolean remove(K key) {
		int idx = indexOf(key);
		Node<K> node = tab[idx];
		if(node != null && (node.key == key || (key!=null && key.equals(node.key)))) {
			tab[idx] = node.next;
			size--;
			return true;
		}
		Node<K> pre = node;
		while(node != null) {
			if(node.key == key || (key!=null && key.equals(node.key))) {
				pre.next = node.next;
				size--;
				return true;
			}
			pre = node;
			node = node.next;
		}
		return false;
	}
	
	public boolean contains(K key) {
		int i = indexOf(key);
		Node node = tab[i];
		while(node != null) {
			if(node.key == key || (key!=null && key.equals(node.key))) {
				return true;
			}
			node = node.next;
		}
		return false;
	}
	
	public int size() {
		return size;
	}

	// test cases
	public static void main (String[] args) {
		HashSet<Integer> set = new HashSet<>(16);
		System.out.println(set.size());
		set.add(null);
		System.out.println(set.contains(null));
		System.out.println(set.size());
		set.add(1);
		set.add(1);
		set.add(2);
		System.out.println(set.contains(2));
		System.out.println(set.size());
		set.remove(1);
		System.out.println(set.contains(1));
		System.out.println(set.size());
	}
}