White Falcon And Tree Problem

white-falcon.py

# TODO: this program too slow, got be some nice solution. :(
# some graph related knownledge

def f(node, x):
    return node[0] * x + node[1]

paths = {}

def cache(fn):
    def _(u, v, *args, **kwargs):
        key = '%d->%d' % (u, v)
        if key not in paths:
            t, val = fn(u, v, *args, **kwargs)
            paths[key] = (t, val)
            paths['%d-%d' % (v, u)] = (t, val[::-1])
        return paths[key]
    return _


def bfs(graph, start, goal):
    stack = [(start, [start])]

    while stack:
        vertex, path = stack.pop(0)
        for v in (graph[vertex] - set(path)):
            if v == goal:
                return path + [v]
            stack.append((v, path + [v]))


@cache
def get_paths(u, v, paths, parent=None, stack=None):
    if stack is None: stack = []
    stack.append(u)
    if u == v:
        return True, stack

    for c in paths[u]:
        if c == parent: continue
        t, stack = get_paths(c, v, paths, parent=u, stack=stack)
        if t: return True, stack
        stack.pop()
    return False, stack

def calc(nodes, path, x):
    while path:
        x = f(nodes[path.pop(0) - 1], x)
    return x % M


M = 10**9 + 7

N = int(raw_input())
nodes = [map(int, raw_input().split()) for i in xrange(N)]

paths = {}
for i in xrange(N-1):
    k, v = map(int, raw_input().split())
    if k not in paths:
        paths[k] = set()
    if v not in paths:
        paths[v] = set()
    paths[k].add(v)
    paths[v].add(k)


Q = int(raw_input())

for i in xrange(Q):
    query = map(int, raw_input().split())
    if query[0] == 1:
        u, v, a, b = query[1:]
        path = bfs(paths, u, v)
        for idx in path:
            nodes[idx-1] = [a, b]

    elif query[0] == 2:
        u, v, x = query[1:]
        path = bfs(paths, u, v)
        print(calc(nodes, path, x))