Distance Code

def encode():
    nnodes = int(input())
    connections = {}

    for i in range(nnodes - 1):
        f, t = list(map(int, input().split()))
        if f in connections:
            connections[f].append(t)
        else:
            connections[f] = [t]

        if t in connections:
            connections[t].append(f)
        else:
            connections[t] = [f]

    root = 1 # The decoder will assume the root is not a leaf, but that doesn't matter when encoding
             # because the root won't matter for isometry

    dists = []
    nodes = []

    def pop_leafs(node, parent, depth, dist):
        for t in connections[node]:
            if t != parent:
                if dist == None:
                    pop_leafs(t, node, depth + 1, None)
                else:
                    pop_leafs(t, node, depth + 1, dist + 1)
                dist = 1

        nodes.append(node)

        if dist != None:
            dists.append(dist)

    pop_leafs(root, 0, 0, None)

    print(" ".join(str(n) for n in nodes))

    return dists


def decode():
    input()
    dists = list(map(int, input().split()))

    killed_nodes = {1}
    current_node = 1
    max_node_idx = 1
    connections_upwards = {}

    for dist in dists:
        connections_before = list(connections_upwards.keys())

        # Go up one
        if current_node in connections_upwards:
            above = connections_upwards[current_node]
        else:
            # Create new node
            max_node_idx += 1
            above = max_node_idx
            connections_upwards[current_node] = above

        down = above

        # Go down dist - 1 times
        for _ in range(dist - 1):
            # Create new node
            max_node_idx += 1
            new = max_node_idx

            connections_upwards[new] = down

            down = new

        current_node = down

        killed_nodes.add(current_node)

        # print("After {} at {}".format(dist, current_node))
        # print("Killed:", killed_nodes)
        # print("New connections:")

        # for f in connections_upwards:
        #     if f not in connections_before:
        #         print("\t", f, "->", connections_upwards[f])

        # print()

    for f in connections_upwards:
        print(f, connections_upwards[f])

t = input()

if t == "1":
    encode()
else:
    decode()

Test 1

Group: 1, 2, 3

Verdict: ACCEPTED

Test 2

Group: 1, 2, 3

Verdict: ACCEPTED

Test 3

Group: 1, 2, 3

Verdict: ACCEPTED

Test 4

Group: 1, 2, 3

Verdict: ACCEPTED

Test 5

Group: 1, 2, 3

Verdict: ACCEPTED

Test 6

Group: 1, 2, 3

Verdict: ACCEPTED

Test 7

Group: 1, 2, 3

Verdict: ACCEPTED

Test 8

Group: 1, 2, 3

Verdict: ACCEPTED

Test 9

Group: 1, 2, 3

Verdict: ACCEPTED

Test 10

Group: 1, 2, 3

Verdict: ACCEPTED

Test 11

Group: 1, 2, 3

Verdict: ACCEPTED

Test 12

Group: 2, 3

Verdict: ACCEPTED

Test 13

Group: 2, 3

Verdict: ACCEPTED

Test 14

Group: 2, 3

Verdict: ACCEPTED

Test 15

Group: 2, 3

Verdict: ACCEPTED

Test 16

Group: 3

Verdict: TIME LIMIT EXCEEDED

Test 17

Group: 3

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 97, in <module>
    encode...

Test 18

Group: 3

Verdict: TIME LIMIT EXCEEDED

Test 19

Group: 3

Verdict: TIME LIMIT EXCEEDED

Test 20

Group: 1, 2, 3

Verdict: ACCEPTED