CSES Solution - Finding Patterns

//code by flutterfly
using System;
using System.Collections.Generic;

public class Program
{
    // Structure for each node in the trie
    public class TrieNode
    {
        public int failureLink;
        public int[] child = new int[26];
        public int occurrences = 0;
        public List<int> patternIndices = new List<int>();
    }

    static string s;
    static int n, nodeIndex = 1;
    static int[] patternFound = new int[500005];
    static List<int>[] trieAdjList = new List<int>[500005];
    static TrieNode[] trie = new TrieNode[500005];

    // Function to insert a pattern into the trie
    static void InsertPattern(string pattern, int index)
    {
        // Start at the root node
        int node = 1;

        // Traverse the pattern character by character
        foreach (char c in pattern)
        {
            int idx = c - 'a';

            // If the current character's child node doesn't exist, create it
            if (trie[node].child[idx] == 0)
                trie[node].child[idx] = ++nodeIndex;

            // Move to the child node
            node = trie[node].child[idx];
        }

        // After traversing the pattern, add the pattern index to the current node
        trie[node].patternIndices.Add(index);
    }

    // Function to build the failure links for the Aho-Corasick algorithm
    static void BuildFailureLinks()
    {
        // Initialize a queue for BFS traversal
        Queue<int> nodeQueue = new Queue<int>();

        // Start from the root node
        int node = 1;

        // Set failure link of root node to itself
        trie[1].failureLink = 1;

        // Push all child nodes of the root to the queue
        for (int i = 0; i < 26; i++)
        {
            if (trie[node].child[i] != 0)
            {
                trie[trie[node].child[i]].failureLink = node;
                nodeQueue.Enqueue(trie[node].child[i]);
            }
            else
            {
                trie[node].child[i] = 1; // If child node doesn't exist, point to root
            }
        }

        // Perform BFS traversal to set failure links
        while (nodeQueue.Count > 0)
        {
            node = nodeQueue.Dequeue();
            for (int i = 0; i < 26; i++)
            {
                if (trie[node].child[i] != 0)
                {
                    trie[trie[node].child[i]].failureLink = trie[trie[node].failureLink].child[i];
                    nodeQueue.Enqueue(trie[node].child[i]);
                }
                else
                {
                    trie[node].child[i] = trie[trie[node].failureLink].child[i];
                }
            }
        }

        // Create adjacency list based on failure links
        for (int i = 0; i <= nodeIndex; i++)
            trieAdjList[i] = new List<int>();

        for (int i = 2; i <= nodeIndex; i++)
            trieAdjList[trie[i].failureLink].Add(i);
    }

    // Function to run the Aho-Corasick algorithm on the main string
    static void RunAhoCorasick(string s)
    {
        // Start from the root node
        int node = 1;

        // Traverse the main string character by character
        foreach (char c in s)
        {
            int idx = c - 'a';

            // Traverse the trie based on characters in the main string
            node = trie[node].child[idx];

            // Increment occurrences at each node
            trie[node].occurrences++;
        }
    }

    // DFS function to propagate the occurrences down the trie
    static int DFS(int u)
    {
        // Initialize total occurrences at current node
        int totalOccurrences = trie[u].occurrences;

        // Traverse child nodes recursively
        foreach (int v in trieAdjList[u])
            totalOccurrences += DFS(v);

        // Propagate occurrences to pattern indices
        foreach (int patternIndex in trie[u].patternIndices)
            patternFound[patternIndex] = totalOccurrences;

        // Return total occurrences
        return totalOccurrences;
    }

    // Main function
    public static void Main(string[] args)
    {
        s = "aybabtu";
        n = 3;
        string[] patterns = { "bab", "abc", "ayba" };

        // Input and insert patterns into trie
        for (int i = 0; i < 500005; i++)
            trie[i] = new TrieNode();

        for (int i = 0; i < n; i++)
            InsertPattern(patterns[i], i);

        // Build failure links for Aho-Corasick algorithm
        BuildFailureLinks();

        // Run Aho-Corasick algorithm on main string
        RunAhoCorasick(s);

        // Propagate occurrences down the trie using DFS
        DFS(1);

        // Output pattern matches
        for (int i = 0; i < n; i++)
            Console.WriteLine((patternFound[i] != 0 ? "YES" : "NO"));
    }
}