Implementing Union-Find and KMP Algorithm in Java

Union-Find (Disjoint Set Union)

The Union-Find data structure supports two primary operations in near O(1) time: checking whether two elements belong to the same set, and merging two disjoint sets into one. It maintains three internal mappings:

  • elementMap: maps user-provided values to internal Element wrappers
  • fatherMap: maps each element to its parent element
  • sizeMap: stores the size of the set for each representative element

Core methods

isSameSet checks if two elements are in the same set by comparing thier representative (root) elements. It first verifies that both element are registered.

union merges two sets. After retrieving the representatives, if they are different, the smaller set is attached to the larger one (by updating the parent pointer of the small representative). The size map is updated accordingly and the entry for the old small representative is removed.

findHead returns the representative element of a given element by traversing parent pointers until reaching a self-pointing node. A stack collects nodes along the path, and after finding the root, all visited nodes are directly attached to the root (path compression) to flatten the tree.

package com.example.unionfind;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

public class UnionFind<V> {
    
    public static class Element<V>{
        public V value;
        public Element(V value) {
            this.value = value;
        }
    }
    
    private HashMap<V, Element<V>> elementMap;
    private HashMap<Element<V>, Element<V>> fatherMap;
    private HashMap<Element<V>, Integer> sizeMap;
    
    public UnionFind(List<V> list) {
        elementMap = new HashMap<>();
        fatherMap = new HashMap<>();
        sizeMap = new HashMap<>();
        for (V v : list) {
            Element<V> element = new Element<>(v);
            elementMap.put(v, element);
            fatherMap.put(element, element);
            sizeMap.put(element, 1);
        }
    }
    
    private Element<V> findHead(Element<V> element) {
        Stack<Element<V>> stack = new Stack<>();
        while (element != fatherMap.get(element)) {
            stack.push(element);
            element = fatherMap.get(element);
        }
        while (!stack.isEmpty()) {
            fatherMap.put(stack.pop(), element);
        }
        return element;
    }
    
    public boolean isSameSet(V v1, V v2) {
        if (!elementMap.containsKey(v1) || !elementMap.containsKey(v2)) {
            return false;
        }
        return findHead(elementMap.get(v1)) == findHead(elementMap.get(v2));
    }
    
    public void union(V v1, V v2) {
        if (elementMap.containsKey(v1) && elementMap.containsKey(v2)) {
            Element<V> root1 = findHead(elementMap.get(v1));
            Element<V> root2 = findHead(elementMap.get(v2));
            if (root1 != root2) {
                Element<V> small = sizeMap.get(root1) < sizeMap.get(root2) ? root1 : root2;
                Element<V> big = (small == root1) ? root2 : root1;
                fatherMap.put(small, big);
                sizeMap.put(big, sizeMap.get(root1) + sizeMap.get(root2));
                sizeMap.remove(small);
            }
        }
    }
}

KMP String Matching Algorithm

The KMP algorithm optimizes pattern matching by precomputing a next array that stores the length of the longest proper prefix that is also a suffix for each prefix of the pattern. When a mismatch occurs, the pattern can shift by more than one position using the information in next, avoiding redundant comparisons.

The next array is built by self-matching the pattern. Two pointers i (current position) and cn (length of current longest prefix-suffix) are used. If characters at cn and i-1 match, next[i] is set to cn+1 and both pointers advance. Otherwise, cn is set to next[cn] (fallback to a shorter prefix-suffix) until a match is found or cn becomes 0.

The main matching loop uses p1 to scan the text and p2 to scan the pattern. On equality, both advance. On mismatch, if p2 is not at the start, it is set to next[p2]; otherwise, p1 advances. A match is found when p2 reaches the end of the pattern.

The following implementation includes a benchmark comparing KMP against Java's built-in indexOf.

package com.example.kmp;

public class KMP {
    
    public static void main(String[] args) {
        benchmark();
    }
    
    public static void benchmark() {
        int trials = 500000;
        int maxLen = 1000;
        long systemTime = 0;
        long kmpTime = 0;
        
        for (int i = 0; i < trials; i++) {
            String text = randomString(maxLen);
            String pattern = randomString(maxLen);
            
            long start = System.currentTimeMillis();
            int res1 = text.indexOf(pattern);
            long end = System.currentTimeMillis();
            systemTime += (end - start);
            
            start = System.currentTimeMillis();
            int res2 = kmpSearch(text, pattern);
            end = System.currentTimeMillis();
            kmpTime += (end - start);
            
            if (res1 != res2) {
                System.out.println("Mismatch: " + res1 + " != " + res2);
                System.out.println("Text: " + text);
                System.out.println("Pattern: " + pattern);
                return;
            }
        }
        System.out.println("All tests passed.");
        System.out.println("indexOf total time: " + systemTime + " ms");
        System.out.println("KMP total time: " + kmpTime + " ms");
    }
    
    private static String randomString(int maxLen) {
        int len = (int)(Math.random() * maxLen) + 1;
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < len; i++) {
            char c = (char)((int)(Math.random() * 26) + 'a');
            sb.append(c);
        }
        return sb.toString();
    }
    
    public static int kmpSearch(String text, String pattern) {
        if (text == null || pattern == null || text.length() < pattern.length()) {
            return -1;
        }
        
        char[] txt = text.toCharArray();
        char[] pat = pattern.toCharArray();
        int[] next = computeNext(pattern);
        
        int i = 0; // index for text
        int j = 0; // index for pattern
        
        while (i < txt.length && j < pat.length) {
            if (txt[i] == pat[j]) {
                i++;
                j++;
            } else if (next[j] == -1) {
                i++;
            } else {
                j = next[j];
            }
        }
        
        return (j == pat.length) ? (i - j) : -1;
    }
    
    private static int[] computeNext(String pattern) {
        if (pattern == null || pattern.length() == 0) {
            return null;
        }
        if (pattern.length() == 1) {
            return new int[]{-1};
        }
        
        int[] next = new int[pattern.length()];
        next[0] = -1;
        next[1] = 0;
        
        char[] p = pattern.toCharArray();
        int cn = 0;   // length of the current longest prefix-suffix
        int pos = 2;  // current position in the pattern
        
        while (pos < pattern.length()) {
            if (p[cn] == p[pos - 1]) {
                next[pos] = ++cn;
                pos++;
            } else if (cn > 0) {
                cn = next[cn];
            } else {
                next[pos] = 0;
                pos++;
            }
        }
        return next;
    }
}

Tags: algorithm Union-Find disjoint-set KMP string-matching

Posted on Mon, 26 Oct 2026 16:44:24 +0000 by doddsey_65