Substrings search (Aho-Corasick algorithm)

Question

I have a code where the goal is to give score to strings based on how much score certain substrings are worth (including repeating substrings).

The way my code (which is based on the Aho-Corasick algorithm) works is:

1. Creating a tree-like data structure where each node represents one prefix of a string (or the full string, along with the data needed to calculate score for this match) from the set of substrings being searched for, with some extra links pointing from some of the nodes a to a node that represents the largest suffix of the string that a represents, to avoid backtracking.
2. The search function uses this structure to find all the substrings (not remembering the matched results, only adding score if a match is found) with near O(N) complexity in most cases.

The performance bottleneck of my code is interlink() which creates these "extra links" in the tree, everything else works reasonably fast, any suggestions for how to improve this function's time complexity or merge it with the tree creation process (while reducing time complexity) are greatly appreciated.

The input strings for createTree() and scan() are expected to contain only lowercase a-z characters, and they can get as long as millions of chars.

import java.util.ArrayList;
import java.util.Arrays;

public class Tree{
    private int[] scores;
    private int[] scoreIndexes;
    private Tree[] branches = new Tree[26];
    private Tree link;

    private Tree(){}
    /**
     * @param strs - The substrings this tree will search for.
     * @param scores - An array of equal length to <b>strs</b> where each <i>scores[i]</i> represents how much score <i>strs[i]</i> is worth.
     */
    public Tree(String[] strs, int[] scores){
        this.scores = scores;
        Tree currentNode;
        String str;
        int strLength;

        for(int i = 0; i < strs.length; i++){
            currentNode = this;
            str = strs[i];
            strLength = str.length();

            for(int c = 0; c < strLength; c++){
                int idx = str.charAt(c) - 'a';
                if(currentNode.branches[idx] == null)
                    currentNode.branches[idx] = new Tree();
                currentNode = currentNode.branches[idx];
                if(strLength == c + 1)
                    currentNode.addScoreIndex(i);
            }
        }
        interlink();
    }

    /**
     * @param str - The string to be scanned.
     * @param minIdx - The first index from which scores from the tree's <b>scores</b> array are counted towards the total value of <b>str</b>.
     * @param maxIdx - The last index from which scores from the tree's <b>scores</b> array are counted towards the total value of <b>str</b>.
     * @return - The score <b>str</b> is worth.
     */
    public long scan(String str, int minIdx, int maxIdx){
        long score = 0L;
        ArrayList<Tree> nodes = new ArrayList<Tree>();
        int branchIndex;
        Tree node;
        for(int index = 0; index < str.length(); index++){
            branchIndex = str.charAt(index) - 'a';

            for(int e = 0; e < nodes.size(); e++) // advance nodes to next char
                nodes.set(e, nodes.get(e).branches[branchIndex]);

            while(nodes.remove(null));
            if(nodes.isEmpty()){ // if no nodes left, continue the search from first node
                nodes.add(this.branches[branchIndex]);
                nodes.remove(null);
            }
            for(int e = 0; e < nodes.size(); e++){
                node = nodes.get(e);
                if(node.link != null)
                    if(!nodes.contains(node.link))
                        nodes.add(node.link);

                if(node.scoreIndexes != null)
                    for(int x = 0; x < node.scoreIndexes.length; x++)
                        if(node.scoreIndexes[x] >= minIdx && node.scoreIndexes[x] <= maxIdx)
                            score += scores[node.scoreIndexes[x]];
            }
        }
        return score;
    }

    /** Creates links in the tree, which are used to avoid backtracking in <b>scan(...)</b>. */
    private void interlink(){
        int[] path = new int[1];
        for(int i = 0; i < 26; i++)
            if(branches[i] != null){
                path[0] = i;
                branches[i].interlinkRecursive(this, path);
            }
    }
    private void interlinkRecursive(Tree mainTree, int[] path){
        int pathLength = path.length;
        int[] newPath = null;
        for(int i = 0; i < 26; i++)
            if(branches[i] != null){
                if(newPath == null)
                    newPath = Arrays.copyOf(path, pathLength + 1);
                newPath[pathLength] = i;
                branches[i].interlinkRecursive(mainTree, newPath);
            }
        if(pathLength < 2) return;
        Tree node;
        for(int k = 1; k < pathLength; k++){
            node = mainTree;
            for(int c = k; c < pathLength; c++){
                node = node.branches[path[c]];
                if(node == null) break;
            }
            if(node != null){
                link = node;
                break;
            }
        }
    }

    private void addScoreIndex(int index){
        if(scoreIndexes == null) scoreIndexes = new int[1];
        else scoreIndexes = Arrays.copyOf(scoreIndexes, scoreIndexes.length + 1);
        scoreIndexes[scoreIndexes.length - 1] = index;
    }
}

Answer 1

Thanks for sharing your code.

Note that what follows are guesses because I can't run your code, but might gives you an idea for better performance.

---

First thing in your interlink method you call path.length() a few times, the Sun implementation is O(1) but it still adds function calls for no good reason (unless the compiler is smart enough to change it itself). So I'd do something like that:

private void interlink(Tree mainTree, String path){
    int pathLength = path.length(); // compute it once here
    for(int i = 0; i < branches.length; i++){
        if(branches[i] != null)
            branches[i].interlink(mainTree, path + (char)(i + 'a'));
    }
    if(pathLength < 2) return;
    Tree node;
    for(int k = 1; k < pathLength; k++){

        node = mainTree;
        for(int c = k; c < pathLength; c++){
            node = node.branches[path.charAt(c) - 'a'];
            if(node == null)
                break;
        }
        if(node != null){
            link = node;
            break;
        }
    }
}

I think it also is a good things to keep in mind in general as function calls are sometimes expensives.

Another thing is you could use StringBuilder to do the String concatenation.

for(int i = 0; i < branches.length; i++){
    if(branches[i] != null)
        builder = new StringBuilder();
        builder.append(path);
        builder.append((char)(i + 'a'));
        branches[i].interlink(mainTree, builder.toString());
}

But the overhead of creating a new StringBuilder inside the loop might kill the performance even more.

Last thing I see is (char)(i + 'a'). Instead of casting the char, if you had an array of char like:

char[] alphabet = {'a', 'b', 'c' ....};

You could use it like

for(int i = 0; i < branches.length; i++){
    if(branches[i] != null)
        branches[i].interlink(mainTree, path + alphabet[i]);
}

(or with the StringBuilder version)

And maybe make it a static member of your class Tree to not have a separate copy of it on all instance of your Tree.

---

All of those above are cheap tricks that probably won't make your program run ten times as fast, the biggest gain would be to call interlink less but I cannot see how to do that right now. Good luck!

Answer 2

I solved my problem, the trick was to create the links row by row in the tree so I can use the links in the higher rows to find the largest suffix faster. Now it's complexity is O(N). Here is the final code: (also includes some other changes I made like adding support for bigger numbers)

import java.math.BigInteger;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.Queue;

public class Tree{
    private int[] scores;
    private int[] scoreIndexes;
    private Tree[] branches = new Tree[26];
    private Tree link, parent, shortcut;

    private Tree(){}
    /**
     * @param strs - The substrings this tree will search for.
     * @param scores - An array of equal length to <b>strs</b> where each <i>scores[i]</i> represents how much score <i>strs[i]</i> is worth.
     */
    public Tree(String[] strs, int[] scores){
        this.scores = scores;
        Tree currentNode;
        String str;
        int strLength;

        for(int i = 0; i < strs.length; i++){
            currentNode = this;
            str = strs[i];
            strLength = str.length();

            for(int c = 0; c < strLength; c++){
                int idx = str.charAt(c) - 'a';
                if(currentNode.branches[idx] == null){
                    currentNode.branches[idx] = new Tree();
                    currentNode.branches[idx].parent = currentNode;
                }
                currentNode = currentNode.branches[idx];
                if(strLength == c + 1)
                    currentNode.addScoreIndex(i);
            }
        }
        interlink();
    }

    /**
     * @param str - The string to be scanned.
     * @param minIdx - The first index from which scores from the tree's <b>scores</b> array are counted towards the total value of <b>str</b>.
     * @param maxIdx - The last index from which scores from the tree's <b>scores</b> array are counted towards the total value of <b>str</b>.
     * @return - The score <b>str</b> is worth.
     */
    public BigInteger scan(String str, int minIdx, int maxIdx){
        BigInteger score = new BigInteger("0");
        //long score = 0;
        Tree currentNode = this;
        Tree nextNode;
        int branchIndex;
        int strLength = str.length();
        for(int c = 0; c < strLength; c++){
            branchIndex = str.charAt(c) - 'a';

            while((nextNode = currentNode.branches[branchIndex]) == null){
                if(currentNode == this){
                    if(++c == strLength) return score;
                    branchIndex = str.charAt(c) - 'a';
                }
                else currentNode = currentNode.link;
            }
            currentNode = nextNode;

            while(nextNode != null){
                score = score.add(new BigInteger(Long.toString(nextNode.getScore(scores, minIdx, maxIdx))));// += nextNode.getScore(scores, minIdx, maxIdx);
                nextNode = nextNode.shortcut;
            }
        }
        return score;
    }

    private void interlink(){

        Queue<Tree> queue = new LinkedList<>();
        queue.add(this);
        Tree node;
        while(!queue.isEmpty()){
            node = queue.poll();
            for(int i = 0; i < branches.length; i++)
                if(node.branches[i] != null){
                    node.branches[i].setLinks(i);
                    queue.add(node.branches[i]);
                }
        }
    }
    private void setLinks(int index){
        Tree currentNode = parent;
        while(link == null){
            if(currentNode.parent == null){
                link = currentNode;
            }
            else{
                currentNode = currentNode.link;
                link = currentNode.branches[index];
            }
        }
        shortcut = link;
        while(shortcut != null){
            if(shortcut.scoreIndexes != null)
                break;
            shortcut = shortcut.link;
        }
    }

    private void addScoreIndex(int index){
        if(scoreIndexes == null) scoreIndexes = new int[1];
        else scoreIndexes = Arrays.copyOf(scoreIndexes, scoreIndexes.length + 1);
        scoreIndexes[scoreIndexes.length - 1] = index;
    }

    private long getScore(int[] scores, int minIdx, int maxIdx){
        long score = 0;
        if(scoreIndexes != null)
            for(int x = 0; x < scoreIndexes.length; x++)
                if(scoreIndexes[x] >= minIdx && scoreIndexes[x] <= maxIdx)
                    score += scores[scoreIndexes[x]];
        return score;
    }
}