CSES Solutions - Substring Order I

import java.util.*;

// Define a Node class representing each node in the suffix
// tree
class Node {
    // Number of distinct substrings in the subtree rooted
    // at this node
    long count;
    // Length of the longest suffix in the subtree rooted at
    // this node
    int length;
    // Suffix link
    int link;
    // Next nodes in the trie structure
    Map<Character, Integer> nextNodes;

    // Constructor to initialize node properties
    Node()
    {
        count = 0;
        length = 0;
        link = -1;
        nextNodes = new HashMap<>();
    }
}

public class SuffixTree {
    static final int MAX_N = 100005;
    // Array to store nodes of the suffix tree
    static Node[] nodes;
    // Variables to track size and last node index
    static int size, last;
    // List to store the kth smallest distinct substring
    static List<Character> result;

    public static void main(String[] args)
    {
        String inputString = "babaacbaab";
        long kthSmallest = 10;

        // Initialize suffix tree
        init();
        for (int i = 0; i < inputString.length(); i++) {
            extend(inputString.charAt(i));
        }
        // Calculate the number of distinct substrings in
        // each subtree
        calculate();

        // DFS traversal to find kth smallest distinct
        // substring
        dfs(0, kthSmallest - 1);

        // Output the kth smallest distinct substring
        for (char c : result) {
            System.out.print(c);
        }
        System.out.println();
    }

    // Initialize the suffix tree
    static void init()
    {
        // Initialize array to store nodes
        nodes = new Node[2 * MAX_N];
        for (int i = 0; i < nodes.length; i++) {
            // Create new node objects
            nodes[i] = new Node();
        }
        // Root node properties
        nodes[0].length = 0;
        nodes[0].link = -1;
        // Initialize size and last node index
        size = 1;
        last = 0;
        // Initialize result list
        result = new ArrayList<>();
    }

    // Extend the suffix tree with a new character
    static void extend(char character)
    {
        int current = size++;
        nodes[current].length = nodes[last].length + 1;
        int previous = last;
        while (previous != -1
               && !nodes[previous].nextNodes.containsKey(
                   character)) {
            nodes[previous].nextNodes.put(character,
                                          current);
            previous = nodes[previous].link;
        }
        if (previous == -1) {
            // If no common suffix, link to root
            nodes[current].link = 0;
        }
        else {
            // Get index of next node
            int next
                = nodes[previous].nextNodes.get(character);
            if (nodes[previous].length + 1
                == nodes[next].length) {
                // If current length is previous length + 1,
                // link to next node
                nodes[current].link = next;
            }
            else {
                // Create a clone of next node
                int clone = size++;
                // Update clone node properties
                nodes[clone].length
                    = nodes[previous].length + 1;
                nodes[clone].nextNodes
                    = new HashMap<>(nodes[next].nextNodes);
                nodes[clone].link = nodes[next].link;
                while (previous != -1
                       && nodes[previous].nextNodes.get(
                              character)
                              == next) {
                    // Update previous node's next nodes to
                    // point to clone
                    nodes[previous].nextNodes.put(character,
                                                  clone);
                    // Move to previous node's suffix link
                    previous = nodes[previous].link;
                }
                // Update current and next node's suffix
                // links
                nodes[next].link = nodes[current].link
                    = clone;
            }
        }
        // Update last node index
        last = current;
    }

    // Calculate the number of distinct substrings in each
    // subtree
    static void calculate()
    {
        // Start with the root node
        calculate(0);
    }

    // Recursive method to calculate the number of distinct
    // substrings in each subtree
    static void calculate(int nodeIndex)
    {
        // Initialize count for current node
        nodes[nodeIndex].count = 1;
        for (Map.Entry<Character, Integer> entry :
             nodes[nodeIndex].nextNodes.entrySet()) {
            // Get index of next node
            int nextNodeIndex = entry.getValue();
            if (nodes[nextNodeIndex].count == 0) {
                // Recursively calculate count for next node
                calculate(nextNodeIndex);
            }
            // Update count for current node
            nodes[nodeIndex].count
                += nodes[nextNodeIndex].count;
        }
    }

    // DFS traversal to find kth smallest distinct substring
    static void dfs(int nodeIndex, long k)
    {
        if (k < 0)
            // Base case: stop traversal if k is negative
            return;
        for (Map.Entry<Character, Integer> entry :
             nodes[nodeIndex].nextNodes.entrySet()) {
            // Get current character
            char character = entry.getKey();
            // Get index of next node
            int nextNodeIndex = entry.getValue();
            if (nodes[nextNodeIndex].count <= k) {
                // Decrement k if count is less than or
                // equal to k
                k -= nodes[nextNodeIndex].count;
            }
            else {
                // Add character to result list
                result.add(character);
                // Recursively traverse to next node
                dfs(nextNodeIndex, k - 1);
                // Stop traversal
                return;
            }
        }
    }
}