3331. Find Subtree Sizes After Changes

SourceBiweekly Contest 142 Q2DifficultyMediumRating2045

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

Thinking

A node may be reattached under its nearest ancestor with the same letter. With \(n \le 10^5\) we should not rebuild the edge list and then recount subtrees.

One DFS is enough: stacks \(d[c]\) store ancestors of letter \(c\). Before returning, we add the current subtree size to the previous same-letter ancestor if it exists, otherwise to the parent.

Post-order ensures \(\textit{ans}[i]\) already includes every descendant; popping the stack restores the ancestor chain.

Python3JavaC++GoTypeScript

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

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);
    }
}

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 = [&](this 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(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(0, -1);
        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
}

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: 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 \(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.

Python3JavaC++GoTypeScript

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

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<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;
    }
};

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(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

Description

Solutions

Solution 1

Solution 2: Explicit Stack

Comments