Tree Permutations

You are given a tree with $N$ nodes rooted at 1. Now you must assign values $1,2,....N$ to every node and each value can be assigned to only one node. Find the number of ways to assign these N values to N nodes such that for every node, its value is the minimum of all values in its subtree. Output the answer modulo ${10}^9+7$.

Example:

Input:  5
1 2
1 3
1 4
3 5
Output: 12

Approach

Java

import java.util.ArrayList;
import java.util.List;

public class TreePermutations {
    static int MOD = 1000000007;
    static long[] fact;
    static long ans;

    public static void main(String args[]) throws Exception {
        int N = 5;

        Node[] nodes = new Node[N];
        for (int n = 0; n < N; n++) {
            nodes[n] = new Node();
        }

        int arr[][] = { { 1, 2 }, { 1, 3 }, { 1, 4 }, { 3, 5 } };

        for (int n = 0; n < N - 1; n++) {
            Node a = nodes[arr[n][0] - 1];
            Node b = nodes[arr[n][1] - 1];
            a.linked.add(b);
            b.linked.add(a);
            a.count.add(0);
            b.count.add(0);
        }

        // dfs(null,nodes[0]);
        dfs(nodes[0], N);

        fact = new long[N + 1];
        fact[0] = 1;
        for (int i = 1; i <= N; i++) {
            fact[i] = (i * fact[i - 1]) % MOD;
        }
        ans = 1;

        for (int n = 0; n < N; n++) {
            nodes[n].visited = false;
        }

        // dfs(null,nodes[0],N);
        dfsAns(nodes[0], N);

        System.out.println(ans);

    }

    static void dfsAns(Node root, int N) {

        Stack stack = new Stack(N);
        Stack traversal = new Stack(N);
        stack.push(root);
        while (!stack.isEmpty()) {
            Node node = stack.peek();
            if (!node.visited) {
                node.visited = true;
                Node parent = traversal.peek();
                if (parent != null) {
                    comb(parent.treeCount, node.treeCount);
                    parent.treeCount -= node.treeCount;
                }
                node.treeCount--;
                traversal.push(node);
                for (int i = 0; i < node.linked.size(); i++) {
                    Node child = node.linked.get(i);
                    if (child.visited)
                        continue;
                    stack.push(child);
                }
            } else {
                traversal.pop();
                stack.pop();
            }
        }

    }

    static void dfs1(Node parent, Node node, int T) {
        // System.out.println(T +" " +node.treeCount);
        comb(T, node.treeCount);
        int nt = node.treeCount - 1;
        for (int i = 0; i < node.linked.size(); i++) {
            Node child = node.linked.get(i);
            if (parent == child)
                continue;
            dfs1(node, child, nt);
            nt -= node.count.get(i);
        }
    }

    static void comb(int n, int r) {
        if (n > 1 && n != r) {
            ans = (ans * fact[n]) % MOD;
            ans = (ans * modInv(fact[r])) % MOD;
            ans = (ans * modInv(fact[n - r])) % MOD;
        }
    }

    static long x, y;

    static long modInv(long a) {
        euclid(a, MOD);
        return (x % MOD + MOD) % MOD;
    }

    static void euclid(long a, long b) {
        if (b == 0) {
            x = 1;
            y = 0;
        } else {
            euclid(b, a % b);
            long t = x;
            x = y;
            y = t - (a / b) * y;
        }
    }

    static void dfs(Node root, int N) {

        Stack stack = new Stack(N);
        Stack traversal = new Stack(N);
        stack.push(root);
        while (!stack.isEmpty()) {
            Node node = stack.peek();
            if (!node.visited) {
                node.visited = true;
                node.treeCount = 1;
                traversal.push(node);
                for (int i = 0; i < node.linked.size(); i++) {
                    Node child = node.linked.get(i);
                    if (child.visited)
                        continue;
                    stack.push(child);
                }
            } else {
                traversal.pop();
                Node parent = traversal.peek();
                if (parent != null) {
                    parent.updateCnt(node);
                }
                stack.pop();
            }
        }

    }

    static int dfs1(Node parent, Node node) {
        node.treeCount = 1;
        for (int i = 0; i < node.linked.size(); i++) {
            Node child = node.linked.get(i);
            if (parent == child)
                continue;
            int chCnt = dfs1(node, child);
            node.treeCount += chCnt;
            node.count.set(i, chCnt);
        }
        return node.treeCount;
    }

    static class Node {
        boolean visited;
        int treeCount;
        List<Node> linked = new ArrayList<>();
        List<Integer> count = new ArrayList<>();

        void updateCnt(Node child) {
            this.treeCount += child.treeCount;
            for (int i = 0; i < linked.size(); i++) {
                if (linked.get(i) == child) {
                    count.set(i, child.treeCount);
                    return;
                }
            }
        }
    }

    static class Stack {
        Node[] data;
        int pointer;

        public Stack(int size) {
            this.data = new Node[size];
            this.pointer = 0;
        }

        void push(Node node) {
            this.data[pointer++] = node;
        }

        Node peek() {
            if (isEmpty())
                return null;
            return this.data[pointer - 1];
        }

        Node pop() {
            return this.data[--pointer];
        }

        boolean isEmpty() {
            return this.pointer == 0;
        }
    }
}