1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90
| ''' Author: NEFU AB-IN Date: 2024-07-27 10:24:54 FilePath: \LeetCode\2836\2836.py LastEditTime: 2024-07-27 11:48:47 '''
import random from collections import Counter, defaultdict, deque from datetime import datetime, timedelta from functools import lru_cache from heapq import heapify, heappop, heappush, nlargest, nsmallest from itertools import combinations, compress, permutations, starmap, tee from math import ceil, comb, fabs, floor, gcd, hypot, log, perm, sqrt from string import ascii_lowercase, ascii_uppercase from sys import exit, setrecursionlimit, stdin from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union
TYPE = TypeVar('TYPE') N = int(2e5 + 10) M = int(20) INF = int(1e12) OFFSET = int(100) MOD = int(1e9 + 7)
setrecursionlimit(int(2e9))
class Arr: array = staticmethod(lambda x=0, size=N: [x() if callable(x) else x for _ in range(size)]) array2d = staticmethod(lambda x=0, rows=N, cols=M: [Arr.array(x, cols) for _ in range(rows)]) graph = staticmethod(lambda size=N: [[] for _ in range(size)])
class Math: max = staticmethod(lambda a, b: a if a > b else b) min = staticmethod(lambda a, b: a if a < b else b)
class IO: input = staticmethod(lambda: stdin.readline().rstrip("\r\n")) read = staticmethod(lambda: map(int, IO.input().split())) read_list = staticmethod(lambda: list(IO.read()))
class Std: class TreeAncestor: """ Binary Lifting for Tree Ancestor Queries, allows us to find the 2^i-th ancestor of any node. Ensure that each node has only one edge pointing to another node to apply binary lifting. """
def __init__(self, n: int, m: int, parent: List[int]): """ Initializes the TreeAncestor with the given number of nodes and parent list.
Args: n (int): Number of nodes. m (int): Maximum power of 2 to consider (default calculated based on n). parent (List[int]): List where parent[i] is the parent of node i. """ self.n = n self.m = m pa = [[(p, p)] + Arr.array((-1, -1), m) for p in parent] for i in range(m): for x in range(n): p, s = pa[x][i] pp, ss = pa[p][i] pa[x][i + 1] = (pp, ss + s) self.pa = pa
class Solution: def getMaxFunctionValue(self, receiver: List[int], k: int) -> int: ta = Std.TreeAncestor(len(receiver), k.bit_length() - 1, receiver) ans = 0 for i in range(ta.n): sum_ = i node = i for j in range(ta.m + 1): if (1 << j) & k: node, s = ta.pa[node][j] sum_ += s ans = Math.max(ans, sum_) return ans
|