Hackerearth - Learning Graph

Question

Problem:

Monk goes to graph city to learn about graphs, and meets an undirected graph having \$ N \$ nodes, where each node has a value \$ i \$ such that \$ 1 \le i \le N \$. Each node of the graph is very curious and wants to know something about the nodes which are directly connected to it.

For each node, if we sort the nodes directly connected to it, according to their values in descending order, what is the value of the node at the \$ k \$ th position? The graph is 1-based.

Note: If two nodes have the same value, they are sorted according to their indices in ascending order.

My accepted solution:

import sys
from collections import defaultdict

node_value_dict = defaultdict(int)
connected_graph_component = defaultdict(list)


def connect_graph():
    n, m, k = list(map(int, sys.stdin.readline().split()))
    node_values = list(map(int, sys.stdin.readline().split()))
    for i in range(1, len(node_values) + 1):
        node_value_dict[i] = node_values[i - 1]

    for m_i in range(m):
        n1, n2 = list(map(int, input().split()))
        connected_graph_component[n1].append(n2)
        connected_graph_component[n2].append(n1)

    for i in range(1, n + 1):
        if i in connected_graph_component:
            curr_co = connected_graph_component[i]
            sort_list = sorted(curr_co, key=lambda value: node_value_dict[
                               value], reverse=True)
            for i in range(len(sort_list) - 1):
                if node_value_dict[sort_list[i]] == node_value_dict[sort_list[i + 1]] and sort_list[i] < sort_list[i + 1]:
                    sort_list[i], sort_list[
                        i + 1] = sort_list[i + 1], sort_list[i]

            print(sort_list[k - 1])
        else:
            print(-1)

connect_graph()

With a runtime of 9.14092 sec.

How can I improve the code quality and further optimize runtime duration?

In case of equal values, I need to sort the nodes according to their indices in ascending order. Currently I'm doing it by looping over the sorted list and then comparing values again side by side maintaining two pointers.

Can you please suggest some elegant way of doing this.

Question Link

Answer 1

As for runtime duration, there is not really much you can do, as the algorithm is already optimum.

However, code can be simplified a bit. Here are some tips:

• node_value_dict has same data as node_values -> not needed
• Use input() instead of sys.stdin.readline().split() so you can avoid an import
• Is there a need for connected_graph_component to be a dict? It should be fine as an array of arrays
• As for the sorting, note that you can make the key a pair instead of only 1 value and it will sort first by first value, and if equal it sorts according to second value.
• list(map(int, sys.stdin.readline().split())) is repeated multiple times, put it in a function
• Sometimes batching print calls in 1 call instead of printing for each value is faster

This is my version of the code after applying these tips:

    def read_line_values():
        return list(map(int, input().strip().split(' ')))

    def connect_graph():
        n, m, k = read_line_values()
        node_values = read_line_values()

        connected_graph_component = [[] for _ in range(n + 1)]

        answer = []

        for m_i in range(m):
            n1, n2 = read_line_values()
            connected_graph_component[n1].append(n2)
            connected_graph_component[n2].append(n1)

        for i in range(1, n + 1):

            neighbours = connected_graph_component[i]
            if len(neighbours) < k:
                answer.append("-1")
            else:  
                neighbours.sort(key = lambda node: (node_values[node - 1], node), reverse=True)
                answer.append(str(neighbours[k - 1]))

        print("\n".join(answer))        
    connect_graph()

On a side note, if you run the same code as python 2.7.6 instead of 3.5.2, it´s about 2s faster.