Red Blue Tree (Codeforces 1208H)
Tags: Implement, trees

You are given a tree of $$$n$$$ nodes. The tree is rooted at node $$$1$$$, which is not considered as a leaf regardless of its degree.Each leaf of the tree has one of the two colors: red or blue. Leaf node $$$v$$$ initially has color $$$s_{v}$$$.The color of each of the internal nodes (including the root) is determined as follows.   Let $$$b$$$ be the number of blue immediate children, and $$$r$$$ be the number of red immediate children of a given vertex.  Then the color of this vertex is blue if and only if $$$b - r \ge k$$$, otherwise red. Integer $$$k$$$ is a parameter that is same for all the nodes.You need to handle the following types of queries:   1 v: print the color of node $$$v$$$;  2 v c: change the color of leaf $$$v$$$ to $$$c$$$ ($$$c = 0$$$ means red, $$$c = 1$$$ means blue);  3 h: update the current value of $$$k$$$ to $$$h$$$. 

Input: The first line of the input consists of two integers $$$n$$$ and $$$k$$$ ($$$2 \le n \le 10^{5}$$$, $$$-n \le k \le n$$$) — the number of nodes and the initial parameter $$$k$$$.Each of the next $$$n - 1$$$ lines contains two integers $$$u$$$ and $$$v$$$ ($$$1 \le u,v \le n$$$), denoting that there is an edge between vertices $$$u$$$ and $$$v$$$.The next line consists of $$$n$$$ space separated integers — the initial array $$$s$$$ ($$$-1 \le s_i \le 1$$$). $$$s_{i} = 0$$$ means that the color of node $$$i$$$ is red. $$$s_{i} = 1$$$ means that the color of node $$$i$$$ is blue. $$$s_{i} = -1$$$ means that the node $$$i$$$ is not a leaf.The next line contains an integer $$$q$$$ ($$$1 \le q \le 10^5$$$), the number of queries.$$$q$$$ lines follow, each containing a query in one of the following queries:   1 v ($$$1 \le v \le n$$$): print the color of node $$$v$$$;  2 v c ($$$1 \le v \le n$$$, $$$c = 0$$$ or $$$c = 1$$$): change the color of leaf $$$v$$$ to $$$c$$$ ($$$c = 0$$$ means red, $$$c = 1$$$ means blue). It is guaranteed that $$$v$$$ is a leaf;  3 h ($$$-n \le h \le n$$$): update the current value of $$$k$$$ to $$$h$$$. 

Output: For each query of the first type, print $$$0$$$ if the color of vertex $$$v$$$ is red, and $$$1$$$ otherwise.

Note:  Figures:(i) The initial tree (ii) The tree after the 3rd query (iii) The tree after the 7th query    

Sample Inputs:
5 2
1 2
1 3
2 4
2 5
-1 -1 0 1 0
9
1 1
1 2
3 -2
1 1
1 2
3 1
2 5 1
1 1
1 2

Sample Outputs:
0
0
1
1
0
1

Solution (Python):
from math import nan as NaN

def find_single_paths(childrens, parent, depths):
    depth = 1
    BFS = [1]
    leaves_depth = []
    while len(BFS) > 0:
        next_BFS = []
        for v in BFS:
            children = childrens[v]
            if len(children) == 1 and len(childrens[children[0]]) == 1:
                nodes_to_optimize = []
                top_of_path = children[0]
                bottom_of_path = childrens[top_of_path][0]
                while len(childrens[bottom_of_path]) == 1:
                    nodes_to_optimize.append(bottom_of_path)
                    bottom_of_path = childrens[bottom_of_path][0]
                for n in nodes_to_optimize:
                    parent[n] = top_of_path
                    childrens[n] = [bottom_of_path]
                    childrens[top_of_path] = [bottom_of_path]
                    parent[bottom_of_path] = top_of_path
            next_BFS.extend(children)
        BFS = next_BFS
        depth += 1
    depths[0] = depth + 1
    rdepth = 0
    n = v
    while True:
        depths[n] = rdepth
        n = parent[n]
        rdepth += 1
        if depths[n] >= rdepth:
            break
    for children in childrens:
        for i in range(len(children) - 1):
            if children[i] > children[len(children) - 1]:
                (children[i], children[len(children) - 1]) = (children[len(children) - 1], children[i])

def get_color(v, k, k_age, k_ages):
    if (v, k) in color_from_node_k:
        return color_from_node_k[v, k]
    if k_ages[v] >= k_age:
        color_from_node_k[v, k] = colors[v]
        return colors[v]
    if k > len(childrens[v]):
        color_from_node_k[v, k] = 0
        return 0
    if k <= -len(childrens[v]):
        color_from_node_k[v, k] = 1
        return 1
    max = len(childrens[v])
    min = -max
    new_color = -1
    for child in childrens[v]:
        child_color = get_color(child, k, k_age, k_ages)
        if child_color == 1:
            min += 2
            if min >= k:
                new_color = 1
                break
            continue
        else:
            max -= 2
            if max < k:
                new_color = 0
                break
    deltas[v] = max if max == min else NaN
    k_ages[v] = k_age
    if new_color != colors[v]:
        colors[v] = new_color
    color_from_node_k[v, k] = new_color
    return new_color

def increment_bluer(v, k_age, k_ages):
    while k_ages[v] >= k_age and deltas[v] != NaN:
        deltas[v] += 2
        new_color = 1 if deltas[v] >= k else 0
        if new_color == colors[v]:
            return
        colors[v] = new_color
        k_ages[v] = k_age
        v = parent[v]

def decrement_redder(v, k_age, k_ages):
    while k_ages[v] >= k_age and deltas[v] != NaN:
        deltas[v] -= 2
        new_color = 1 if deltas[v] >= k else 0
        if new_color == colors[v]:
            return
        colors[v] = new_color
        k_ages[v] = k_age
        v = parent[v]

def set_childrens_parent(v, childrens):
    try:
        node_stack = [v]
        while len(node_stack) > 0:
            v = node_stack.pop()
            for child in childrens[v]:
                parent[child] = v
                childrens[child].remove(v)
                node_stack.append(child)
    except KeyError as e:
        print('KeyError ', str(e), 'occurred.')
    except Exception as e:
        print(f'{type(e).__name__}: [{str(e)}] on {v}')
        exit(v)
(n, k) = [int(x) for x in input().split(' ')]
k_age = 0
k_ages = [0] * (n + 1)
parent = [0] * (n + 1)
deltas = [0] * (n + 1)
depths = [0] * (n + 1)
parent[1] = 1
k_ages[0] = -1
childrens_set = [set() for _ in range(n + 1)]
color_from_node_k = {}
for i in range(n - 1):
    (u, v) = [int(x) for x in input().split(' ')]
    if parent[u] != 0:
        parent[v] = u
        childrens_set[u].add(v)
        for child in childrens_set[v]:
            parent[child] = v
            childrens_set[child].remove(v)
            set_childrens_parent(child, childrens_set)
    elif parent[v] != 0:
        parent[u] = v
        childrens_set[v].add(u)
        for child in childrens_set[u]:
            parent[child] = u
            childrens_set[child].remove(u)
            set_childrens_parent(child, childrens_set)
    else:
        childrens_set[u].add(v)
        childrens_set[v].add(u)
parent[1] = 0
childrens = [list(x) for x in childrens_set]
colors = [int(x) for x in input().split(' ')]
colors.insert(0, 0)
find_single_paths(childrens, parent, depths)
k_age += 1
for i in list(range(n, 0, -1)):
    if colors[i] != -1:
        k_ages[i] = float('inf')
q = int(input())
for i in range(q):
    command = [int(x) for x in input().split(' ')]
    if command[0] == 1:
        v = command[1]
        color = get_color(v, k, k_age, k_ages)
        print(color)
        continue
    if command[0] == 2:
        v = command[1]
        if colors[v] != command[2]:
            color_from_node_k = {}
            colors[v] = command[2]
            if command[2] == 1:
                increment_bluer(parent[v], k_age, k_ages)
                continue
            decrement_redder(parent[v], k_age, k_ages)
        continue
    h = command[1]
    k = h
    k_age += 1

Solution (Python):
from math import nan as NaN

def find_single_paths(childrens, parent, depths):
    depth = 1
    BFS = [1]
    leaves_depth = []
    try:
        while len(BFS) > 0:
            next_BFS = []
            for v in BFS:
                children = childrens[v]
                if len(children) == 1 and len(childrens[children[0]]) == 1:
                    nodes_to_optimize = []
                    top_of_path = children[0]
                    bottom_of_path = childrens[top_of_path][0]
                    while len(childrens[bottom_of_path]) == 1:
                        nodes_to_optimize.append(bottom_of_path)
                        bottom_of_path = childrens[bottom_of_path][0]
                    for n in nodes_to_optimize:
                        parent[n] = top_of_path
                        childrens[n] = [bottom_of_path]
                        childrens[top_of_path] = [bottom_of_path]
                        parent[bottom_of_path] = top_of_path
                next_BFS.extend(children)
                if len(children) == 0:
                    depths[v] = depth
                    leaves_depth.append(v)
                    continue
            BFS = next_BFS
            depth += 1
        ideepest = leaves_depth[-1]
        depths[0] = depths[ideepest] + 1
        rdepth = 1
        n = ideepest
        while True:
            depths[n] = rdepth
            n = parent[n]
            rdepth += 1
            if depths[n] >= rdepth:
                break
        for children in childrens:
            for i in range(len(children) - 1):
                if children[i] > children[len(children) - 1]:
                    (children[i], children[len(children) - 1]) = (children[len(children) - 1], children[i])
                if children[i] < children[0]:
                    (children[i], children[0]) = (children[0], children[i])
    except BaseException as e:
        print('Base Exception thrown:', str(e), "occurred. It's a", type(e))
        print('depth v', depth, v)
        print('childrens[v]', end='')
        print(childrens[v])

def get_color(v, k, k_age, k_ages):
    if (v, k) in color_from_node_k:
        return color_from_node_k[v, k]
    if k_ages[v] >= k_age:
        color_from_node_k[v, k] = colors[v]
        return colors[v]
    if k > len(childrens[v]):
        color_from_node_k[v, k] = 0
        return 0
    if k <= -len(childrens[v]):
        color_from_node_k[v, k] = 1
        return 1
    max = len(childrens[v])
    min = -max
    new_color = -1
    for child in childrens[v]:
        child_color = get_color(child, k, k_age, k_ages)
        if child_color == 1:
            min += 2
            if min >= k:
                new_color = 1
                break
            continue
        else:
            max -= 2
            if max < k:
                new_color = 0
                break
    deltas[v] = max if max == min else NaN
    if new_color == -1:
        print('Something totally broke.  new_color == -1. v', v)
        exit(v)
    k_ages[v] = k_age
    if new_color != colors[v]:
        colors[v] = new_color
    color_from_node_k[v, k] = new_color
    return new_color

def increment_bluer(v, k_age, k_ages):
    while k_ages[v] >= k_age and deltas[v] != NaN:
        deltas[v] += 2
        new_color = 1 if deltas[v] >= k else 0
        if new_color == colors[v]:
            return
        colors[v] = new_color
        k_ages[v] = k_age
        v = parent[v]

def decrement_redder(v, k_age, k_ages):
    while k_ages[v] >= k_age and deltas[v] != NaN:
        deltas[v] -= 2
        new_color = 1 if deltas[v] >= k else 0
        if new_color == colors[v]:
            return
        colors[v] = new_color
        k_ages[v] = k_age
        v = parent[v]

def set_childrens_parent(v, childrens):
    try:
        node_stack = [v]
        while len(node_stack) > 0:
            v = node_stack.pop()
            for child in childrens[v]:
                parent[child] = v
                childrens[child].remove(v)
                node_stack.append(child)
    except KeyError as e:
        print('KeyError ', str(e), 'occurred.')
    except Exception as e:
        print(f'{type(e).__name__}: [{str(e)}] on {v}')
        exit(v)
(n, k) = [int(x) for x in input().split(' ')]
k100_0 = n == 100000 and k == 0
k_age = 0
k_ages = [0] * (n + 1)
parent = [0] * (n + 1)
deltas = [0] * (n + 1)
depths = [0] * (n + 1)
parent[1] = 1
k_ages[0] = -1
childrens_set = [set() for _ in range(n + 1)]
color_from_node_k = {}
i = -1
try:
    for i in range(n - 1):
        (u, v) = [int(x) for x in input().split(' ')]
        if parent[u] != 0:
            parent[v] = u
            childrens_set[u].add(v)
            for child in childrens_set[v]:
                parent[child] = v
                childrens_set[child].remove(v)
                set_childrens_parent(child, childrens_set)
        elif parent[v] != 0:
            parent[u] = v
            childrens_set[v].add(u)
            for child in childrens_set[u]:
                parent[child] = u
                childrens_set[child].remove(u)
                set_childrens_parent(child, childrens_set)
        else:
            childrens_set[u].add(v)
            childrens_set[v].add(u)
    parent[1] = 0
    childrens = [list(x) for x in childrens_set]
except Exception as e:
    print(f'{type(e).__name__}: [{str(e)}] on i=={i}')
    exit(i)
debug = False and k100_0 and (parent[9029] == 97016 or parent[97016] == 9029)
if debug:
    print('.', end='')
colors = [int(x) for x in input().split(' ')]
colors.insert(0, 0)
find_single_paths(childrens, parent, depths)
k_age += 1
i = -1
try:
    for i in list(range(n, 0, -1)):
        if colors[i] != -1:
            k_ages[i] = float('inf')
except Exception as e:
    print(f'setting leafs {type(e).__name__}: [{str(e)}] on i=={i}')
    exit(i)
debug_str = ''
known_commands = [[1, 1], [1, 48175], [3, -1], [3, 0], [3, 1]]
q = int(input())
for i in range(q):
    try:
        command = [int(x) for x in input().split(' ')]
        if command[0] == 1:
            v = command[1]
            if debug:
                if (v, k) in color_from_node_k:
                    print('H', end='')
                else:
                    print('-', end='')
            color = get_color(v, k, k_age, k_ages)
            if not debug:
                print(color)
            continue
        if command[0] == 2:
            v = command[1]
            if colors[v] != command[2]:
                color_from_node_k = {}
                colors[v] = command[2]
                if command[2] == 1:
                    increment_bluer(parent[v], k_age, k_ages)
                    continue
                decrement_redder(parent[v], k_age, k_ages)
            continue
        h = command[1]
        k = h
        k_age += 1
    except BaseException as e:
        debug_str += 'Base Exception [' + str(e) + '] '
        debug_str += '[' + str(len(debug_str)) + ']'
        print(debug_str[-400:])

Solution (Python):
from math import nan as NaN

def find_single_paths(childrens, parent, depths):
    depth = 1
    BFS = [1]
    leaves_depth = []
    try:
        while len(BFS) > 0:
            next_BFS = []
            for v in BFS:
                children = childrens[v]
                if len(children) == 1 and len(childrens[children[0]]) == 1:
                    nodes_to_optimize = []
                    top_of_path = children[0]
                    bottom_of_path = childrens[top_of_path][0]
                    while len(childrens[bottom_of_path]) == 1:
                        nodes_to_optimize.append(bottom_of_path)
                        bottom_of_path = childrens[bottom_of_path][0]
                    for n in nodes_to_optimize:
                        parent[n] = top_of_path
                        childrens[n] = [bottom_of_path]
                        childrens[top_of_path] = [bottom_of_path]
                        parent[bottom_of_path] = top_of_path
                next_BFS.extend(children)
                if len(children) == 0:
                    depths[v] = depth
                    leaves_depth.append(v)
                    continue
            BFS = next_BFS
            depth += 1
        ideepest = leaves_depth[-1]
        depths[0] = depths[ideepest] + 1
        rdepth = 1
        n = ideepest
        while True:
            depths[n] = rdepth
            n = parent[n]
            rdepth += 1
            if depths[n] >= rdepth:
                break
        for children in childrens:
            for i in range(len(children) - 1):
                if children[i] > children[len(children) - 1]:
                    (children[i], children[len(children) - 1]) = (children[len(children) - 1], children[i])
    except BaseException as e:
        print('Base Exception thrown:', str(e), "occurred. It's a", type(e))
        print('depth v', depth, v)
        print('childrens[v]', end='')
        print(childrens[v])

def get_color(v, k, k_age, k_ages):
    if (v, k) in color_from_node_k:
        return color_from_node_k[v, k]
    if k_ages[v] >= k_age:
        color_from_node_k[v, k] = colors[v]
        return colors[v]
    if k > len(childrens[v]):
        color_from_node_k[v, k] = 0
        return 0
    if k <= -len(childrens[v]):
        color_from_node_k[v, k] = 1
        return 1
    max = len(childrens[v])
    min = -max
    new_color = -1
    for child in childrens[v]:
        child_color = get_color(child, k, k_age, k_ages)
        if child_color == 1:
            min += 2
            if min >= k:
                new_color = 1
                break
            continue
        else:
            max -= 2
            if max < k:
                new_color = 0
                break
    deltas[v] = max if max == min else NaN
    if new_color == -1:
        print('Something totally broke.  new_color == -1. v', v)
        exit(v)
    k_ages[v] = k_age
    if new_color != colors[v]:
        colors[v] = new_color
    color_from_node_k[v, k] = new_color
    return new_color

def increment_bluer(v, k_age, k_ages):
    while k_ages[v] >= k_age and deltas[v] != NaN:
        deltas[v] += 2
        new_color = 1 if deltas[v] >= k else 0
        if new_color == colors[v]:
            return
        colors[v] = new_color
        k_ages[v] = k_age
        v = parent[v]

def decrement_redder(v, k_age, k_ages):
    while k_ages[v] >= k_age and deltas[v] != NaN:
        deltas[v] -= 2
        new_color = 1 if deltas[v] >= k else 0
        if new_color == colors[v]:
            return
        colors[v] = new_color
        k_ages[v] = k_age
        v = parent[v]

def set_childrens_parent(v, childrens):
    try:
        node_stack = [v]
        while len(node_stack) > 0:
            v = node_stack.pop()
            for child in childrens[v]:
                parent[child] = v
                childrens[child].remove(v)
                node_stack.append(child)
    except KeyError as e:
        print('KeyError ', str(e), 'occurred.')
    except Exception as e:
        print(f'{type(e).__name__}: [{str(e)}] on {v}')
        exit(v)
(n, k) = [int(x) for x in input().split(' ')]
k100_0 = n == 100000 and k == 0
k_age = 0
k_ages = [0] * (n + 1)
parent = [0] * (n + 1)
deltas = [0] * (n + 1)
depths = [0] * (n + 1)
parent[1] = 1
k_ages[0] = -1
childrens_set = [set() for _ in range(n + 1)]
color_from_node_k = {}
i = -1
try:
    for i in range(n - 1):
        (u, v) = [int(x) for x in input().split(' ')]
        if parent[u] != 0:
            parent[v] = u
            childrens_set[u].add(v)
            for child in childrens_set[v]:
                parent[child] = v
                childrens_set[child].remove(v)
                set_childrens_parent(child, childrens_set)
        elif parent[v] != 0:
            parent[u] = v
            childrens_set[v].add(u)
            for child in childrens_set[u]:
                parent[child] = u
                childrens_set[child].remove(u)
                set_childrens_parent(child, childrens_set)
        else:
            childrens_set[u].add(v)
            childrens_set[v].add(u)
    parent[1] = 0
    childrens = [list(x) for x in childrens_set]
except Exception as e:
    print(f'{type(e).__name__}: [{str(e)}] on i=={i}')
    exit(i)
debug = False and k100_0 and (parent[9029] == 97016 or parent[97016] == 9029)
if debug:
    print('.', end='')
colors = [int(x) for x in input().split(' ')]
colors.insert(0, 0)
find_single_paths(childrens, parent, depths)
k_age += 1
i = -1
try:
    for i in list(range(n, 0, -1)):
        if colors[i] != -1:
            k_ages[i] = float('inf')
except Exception as e:
    print(f'setting leafs {type(e).__name__}: [{str(e)}] on i=={i}')
    exit(i)
debug_str = ''
known_commands = [[1, 1], [1, 48175], [3, -1], [3, 0], [3, 1]]
q = int(input())
for i in range(q):
    try:
        command = [int(x) for x in input().split(' ')]
        if command[0] == 1:
            v = command[1]
            if debug:
                if (v, k) in color_from_node_k:
                    print('H', end='')
                else:
                    print('-', end='')
            color = get_color(v, k, k_age, k_ages)
            if not debug:
                print(color)
            continue
        if command[0] == 2:
            v = command[1]
            if colors[v] != command[2]:
                color_from_node_k = {}
                colors[v] = command[2]
                if command[2] == 1:
                    increment_bluer(parent[v], k_age, k_ages)
                    continue
                decrement_redder(parent[v], k_age, k_ages)
            continue
        h = command[1]
        k = h
        k_age += 1
    except BaseException as e:
        debug_str += 'Base Exception [' + str(e) + '] '
        debug_str += '[' + str(len(debug_str)) + ']'
        print(debug_str[-400:])