3331. Find Subtree Sizes After Changes

Description

You are given a tree rooted at node 0 that consists of n nodes numbered from 0 to n - 1. The tree is represented by an array parent of size n, where parent[i] is the parent of node i. Since node 0 is the root, parent[0] == -1.

You are also given a string s of length n, where s[i] is the character assigned to node i.

We make the following changes on the tree one time simultaneously for all nodes x from 1 to n - 1:

Find the closest node y to node x such that y is an ancestor of x, and s[x] == s[y].
If node y does not exist, do nothing.
Otherwise, remove the edge between x and its current parent and make node y the new parent of x by adding an edge between them.

Return an array answer of size n where answer[i] is the size of the subtree rooted at node i in the final tree.

Example 1:

Input: parent = [-1,0,0,1,1,1], s = "abaabc"

Output: [6,3,1,1,1,1]

Explanation:

The parent of node 3 will change from node 1 to node 0.

Example 2:

Input: parent = [-1,0,4,0,1], s = "abbba"

Output: [5,2,1,1,1]

Explanation:

The following changes will happen at the same time:

The parent of node 4 will change from node 1 to node 0.
The parent of node 2 will change from node 4 to node 1.

Constraints:

n == parent.length == s.length
1 <= n <= 10⁵
0 <= parent[i] <= n - 1 for all i >= 1.
parent[0] == -1
parent represents a valid tree.
s consists only of lowercase English letters.

Solutions

Solution 1

Solution 2: Explicit Stack

Thinking

A node may be reattached under its nearest ancestor with the same letter. $n$ can reach $10^5$, so rebuilding the edge list and then recounting subtrees is unnecessary. On a chain the first recursive call always follows the only child, so the depth is $n$ and exceeds the default recursion limit.

A subtree size can be added to its new parent only after every descendant has been counted, and a same-letter ancestor lies on the current path.

An explicit stack walks the tree in postorder, while $\textit{d}[c]$ stores the nodes of letter $c$ still on the path. On entry, set the node’s size to $1$, push it onto its letter stack, then push the exit marker and the children. On exit, add the subtree size to the earlier same-letter node when one exists, otherwise to the parent, and pop the letter stack.

class Solution {
    private List<Integer>[] g;
    private List<Integer>[] d;
    private char[] s;
    private int[] ans;

    public int[] findSubtreeSizes(int[] parent, String s) {
        int n = s.length();
        g = new List[n];
        d = new List[26];
        this.s = s.toCharArray();
        Arrays.setAll(g, k -> new ArrayList<>());
        Arrays.setAll(d, k -> new ArrayList<>());
        for (int i = 1; i < n; ++i) {
            g[parent[i]].add(i);
        }
        ans = new int[n];
        dfs(0, -1);
        return ans;
    }

    private void dfs(int i, int fa) {
        ans[i] = 1;
        int idx = s[i] - 'a';
        d[idx].add(i);
        for (int j : g[i]) {
            dfs(j, i);
        }
        int k = d[idx].size() > 1 ? d[idx].get(d[idx].size() - 2) : fa;
        if (k >= 0) {
            ans[k] += ans[i];
        }
        d[idx].remove(d[idx].size() - 1);
    }
}



// Solution 2
class Solution {
    public int[] findSubtreeSizes(int[] parent, String s) {
        int n = s.length();
        List<Integer>[] g = new List[n];
        List<Integer>[] d = new List[26];
        Arrays.setAll(g, k -> new ArrayList<>());
        Arrays.setAll(d, k -> new ArrayList<>());
        for (int i = 1; i < n; ++i) {
            g[parent[i]].add(i);
        }
        int[] ans = new int[n];
        char[] cs = s.toCharArray();
        Deque<int[]> stk = new ArrayDeque<>();
        stk.push(new int[] {0, -1, 0});
        while (!stk.isEmpty()) {
            int[] cur = stk.pop();
            int i = cur[0], fa = cur[1], state = cur[2];
            int idx = cs[i] - 'a';
            if (state == 0) {
                ans[i] = 1;
                d[idx].add(i);
                stk.push(new int[] {i, fa, 1});
                for (int j : g[i]) {
                    stk.push(new int[] {j, i, 0});
                }
            } else {
                int k = d[idx].size() > 1 ? d[idx].get(d[idx].size() - 2) : fa;
                if (k >= 0) {
                    ans[k] += ans[i];
                }
                d[idx].remove(d[idx].size() - 1);
            }
        }
        return ans;
    }
}

class Solution {
public:
    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        int n = s.size();
        vector<int> g[n];
        vector<int> d[26];
        for (int i = 1; i < n; ++i) {
            g[parent[i]].push_back(i);
        }
        vector<int> ans(n);
        auto dfs = [&](auto&& dfs, int i, int fa) -> void {
            ans[i] = 1;
            int idx = s[i] - 'a';
            d[idx].push_back(i);
            for (int j : g[i]) {
                dfs(dfs, j, i);
            }
            int k = d[idx].size() > 1 ? d[idx][d[idx].size() - 2] : fa;
            if (k >= 0) {
                ans[k] += ans[i];
            }
            d[idx].pop_back();
        };
        dfs(dfs, 0, -1);
        return ans;
    }
};



// Solution 2
class Solution {
public:
    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        int n = s.size();
        vector<vector<int>> g(n);
        vector<int> d[26];
        for (int i = 1; i < n; ++i) {
            g[parent[i]].push_back(i);
        }
        vector<int> ans(n);
        vector<array<int, 3>> stk{{0, -1, 0}};
        while (!stk.empty()) {
            auto cur = stk.back();
            stk.pop_back();
            int i = cur[0], fa = cur[1], state = cur[2];
            int idx = s[i] - 'a';
            if (state == 0) {
                ans[i] = 1;
                d[idx].push_back(i);
                stk.push_back({i, fa, 1});
                for (int j : g[i]) {
                    stk.push_back({j, i, 0});
                }
            } else {
                int k = d[idx].size() > 1 ? d[idx][d[idx].size() - 2] : fa;
                if (k >= 0) {
                    ans[k] += ans[i];
                }
                d[idx].pop_back();
            }
        }
        return ans;
    }
};

class Solution:
    def findSubtreeSizes(self, parent: List[int], s: str) -> List[int]:
        def dfs(i: int, fa: int):
            ans[i] = 1
            d[s[i]].append(i)
            for j in g[i]:
                dfs(j, i)
            k = fa
            if len(d[s[i]]) > 1:
                k = d[s[i]][-2]
            if k != -1:
                ans[k] += ans[i]
            d[s[i]].pop()

        n = len(s)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[parent[i]].append(i)
        d = defaultdict(list)
        ans = [0] * n
        dfs(0, -1)
        return ans



# Solution 2
class Solution:
    def findSubtreeSizes(self, parent: List[int], s: str) -> List[int]:
        n = len(s)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[parent[i]].append(i)
        d = [[] for _ in range(26)]
        ans = [0] * n
        stk = [(0, -1, 0)]
        while stk:
            i, fa, state = stk.pop()
            if state == 0:
                ans[i] = 1
                idx = ord(s[i]) - 97
                d[idx].append(i)
                stk.append((i, fa, 1))
                for j in g[i]:
                    stk.append((j, i, 0))
            else:
                idx = ord(s[i]) - 97
                k = d[idx][-2] if len(d[idx]) > 1 else fa
                if k != -1:
                    ans[k] += ans[i]
                d[idx].pop()
        return ans

func findSubtreeSizes(parent []int, s string) []int {
	n := len(s)
	g := make([][]int, n)
	for i := 1; i < n; i++ {
		g[parent[i]] = append(g[parent[i]], i)
	}
	d := [26][]int{}
	ans := make([]int, n)
	var dfs func(int, int)
	dfs = func(i, fa int) {
		ans[i] = 1
		idx := int(s[i] - 'a')
		d[idx] = append(d[idx], i)
		for _, j := range g[i] {
			dfs(j, i)
		}
		k := fa
		if len(d[idx]) > 1 {
			k = d[idx][len(d[idx])-2]
		}
		if k != -1 {
			ans[k] += ans[i]
		}
		d[idx] = d[idx][:len(d[idx])-1]
	}
	dfs(0, -1)
	return ans
}



// Solution 2
func findSubtreeSizes(parent []int, s string) []int {
	n := len(s)
	g := make([][]int, n)
	for i := 1; i < n; i++ {
		g[parent[i]] = append(g[parent[i]], i)
	}
	d := [26][]int{}
	ans := make([]int, n)
	stk := [][3]int{{0, -1, 0}}
	for len(stk) > 0 {
		cur := stk[len(stk)-1]
		stk = stk[:len(stk)-1]
		i, fa, state := cur[0], cur[1], cur[2]
		idx := int(s[i] - 'a')
		if state == 0 {
			ans[i] = 1
			d[idx] = append(d[idx], i)
			stk = append(stk, [3]int{i, fa, 1})
			for _, j := range g[i] {
				stk = append(stk, [3]int{j, i, 0})
			}
		} else {
			k := fa
			if len(d[idx]) > 1 {
				k = d[idx][len(d[idx])-2]
			}
			if k != -1 {
				ans[k] += ans[i]
			}
			d[idx] = d[idx][:len(d[idx])-1]
		}
	}
	return ans
}

function findSubtreeSizes(parent: number[], s: string): number[] {
    const n = parent.length;
    const g: number[][] = Array.from({ length: n }, () => []);
    const d: number[][] = Array.from({ length: 26 }, () => []);
    for (let i = 1; i < n; ++i) {
        g[parent[i]].push(i);
    }
    const ans: number[] = Array(n).fill(1);
    const dfs = (i: number, fa: number): void => {
        const idx = s.charCodeAt(i) - 97;
        d[idx].push(i);
        for (const j of g[i]) {
            dfs(j, i);
        }
        const k = d[idx].length > 1 ? d[idx].at(-2)! : fa;
        if (k >= 0) {
            ans[k] += ans[i];
        }
        d[idx].pop();
    };
    dfs(0, -1);
    return ans;
}



// Solution 2
function findSubtreeSizes(parent: number[], s: string): number[] {
    const n = parent.length;
    const g: number[][] = Array.from({ length: n }, () => []);
    const d: number[][] = Array.from({ length: 26 }, () => []);
    for (let i = 1; i < n; ++i) {
        g[parent[i]].push(i);
    }
    const ans: number[] = Array(n).fill(0);
    const stk: number[][] = [[0, -1, 0]];
    while (stk.length) {
        const [i, fa, state] = stk.pop()!;
        const idx = s.charCodeAt(i) - 97;
        if (state === 0) {
            ans[i] = 1;
            d[idx].push(i);
            stk.push([i, fa, 1]);
            for (const j of g[i]) {
                stk.push([j, i, 0]);
            }
        } else {
            const k = d[idx].length > 1 ? d[idx][d[idx].length - 2] : fa;
            if (k >= 0) {
                ans[k] += ans[i];
            }
            d[idx].pop();
        }
    }
    return ans;
}

3331 - Find Subtree Sizes After Changes

3331. Find Subtree Sizes After Changes

Description

Solutions

Solution 1

Solution 2: Explicit Stack

All Problems

All Solutions

Welcome to Subscribe On Youtube

3331. Find Subtree Sizes After Changes

Description

Solutions

Solution 1

Solution 2: Explicit Stack

All Problems

All Solutions