MWST Solution

package weblab;

import java.util.*;

class Solution {

    private static List<MWSTNode> getNodesAtDepth(MWSTNode start, int depth) {
        List<MWSTNode> result = new ArrayList<>();
        getNodesAtDepthHelper(result, start, 0, depth);
        return result;
    }

    private static void getNodesAtDepthHelper(List<MWSTNode> result, MWSTNode curNode, int curDepth, int targetDepth) {
        if (curNode == null) return;

        if (curDepth == targetDepth) {
            result.add(curNode);
            return;
        }

        for (MWSTNode c : curNode.getChildren()) {
            getNodesAtDepthHelper(result, c, curDepth+1, targetDepth);
        }
    }

    private static int indexOfChild(MWSTNode parent, MWSTNode child) {
        List<MWSTNode> children = parent.getChildren();

        int i = 0;
        for (MWSTNode c : children) {
            if (Objects.equals(c, child)) return i;
            i++;
        }

        return -1;
    }

    /**
    * See the description of the exercise.
    */
    public static MWSTNode getCousin(MultiWaySearchTree tree, MWSTNode node) {
        if (tree == null || node == null) return null;

        MWSTNode parent = node.getParent();
        if (parent == null) return null;
        MWSTNode grandparent = parent.getParent();
        if (grandparent == null) return null;

        int index = indexOfChild(grandparent, parent) - 1;

        while (index >= 0) {
            MWSTNode start = grandparent.getChildren().get(index);
            List<MWSTNode> cousins = getNodesAtDepth(start, 1);
            if (cousins.size() > 0) return cousins.get(cousins.size()-1);
            index--;
        }

        return null;
    }

    /**
    * See the description of the exercise.
    */
    public static MWSTNode getUncle(MultiWaySearchTree tree, MWSTNode node) {
        if (tree == null || node == null) return null;

        MWSTNode parent = node.getParent();
        if (parent == null) return null;
        MWSTNode grandparent = parent.getParent();
        if (grandparent == null) return null;

        int minIndex = indexOfChild(grandparent, parent) + 1;
        int index = grandparent.getChildren().size()-1;

        while (index >= minIndex) {
            MWSTNode start = grandparent.getChildren().get(index);
            List<MWSTNode> siblings = getNodesAtDepth(start, 0);
            if (siblings.size() > 0) return siblings.get(siblings.size()-1);
            index--;
        }

        return null;
    }

    /**
    * See the description of the exercise.
    */
    public static boolean restrictedSearch(MultiWaySearchTree tree, int key) {
        if (tree == null) return false;
        return searchHelper(tree.getRoot(), 0, 2, key);
    }

    private static boolean searchHelper(MWSTNode curNode, int curDepth, int maxDepth, int key) {
        if (curNode == null) return false;
        if (curDepth > maxDepth) return false;

        List<Integer> keys = curNode.getKeys();
        List<MWSTNode> children = curNode.getChildren();

        for (Integer k : keys)
            if (Objects.equals(k, key)) return true;

        int index = 0;
        for (Integer k : keys) {
            if (key < k) break;
            index++;
        }

        return (searchHelper(children.get(index), curDepth+1, maxDepth, key));
    }
}