I'm trying to run a python data structures code but it's throwing an error

The question is

You are given an undirected graph consisting of N vertices, numbered from 1 to N, and M edges.

The graph is described by two arrays, A and B, both of length M. A pair (A[K], B[K]), for K from 0 to M-1, describes an edge between vertex A[K] and vertex B[K].

Your task is to check whether the given graph contains a path from vertex 1 to vertex N going through all of the vertices, one by one, in increasing order of their numbers. All connections on the path should be direct.

Write a function:

def solution (N, A, B)

that, given an integer N and two arrays A and B of M integers each, returns True if there exists a path from vertex 1 to N going through all vertices, one by one, in increasing order, or False otherwise.

Examples:

  1. Given N = 4, A = [1, 2, 4, 4, 3] and B = [2, 3, 1, 3, 1], the function should return True. There is a path (1-2-3-4) using edges (1, 2), (2, 3) and (4, 3).

  2. Given N = 4, A = [1, 2, 1, 3] and B = [2, 4, 3, 4], the function should return False. There is no path (123→4), as there is no direct connection from vertex 2 to vertex 3.

  3. Given N = 6, A = [2, 4, 5, 3] and B = [3, 5, 6, 4], the function should return False. There is no direct connection from vertex 1 to vertex 2.

  4. Given N = 3, A = [1, 3] and B = [2, 2], the function should return True. There is a path (1-2-3) using edges (1, 2) and (3, 2).

Write an efficient algorithm for the following assumptions:

• N is an integer within the range [2..100,000];

M is an integer within the range [0..100,000];

• all elements of arrays A and B are integers within the range [1..N];

there are no self-loops (edges with A[K] = B[K]) in the graph;

• there are no multiple edges between the same vertices

And the code i've written is :

def solution(N, A, B):
    # write your code in Python 3.6
    pass
    graph = {}
    for i in range(len(A)):
        if A[i] not in graph:
            graph[A[i]] = [B[i]]
        else:
            graph[A[i]].append(B[i])
        if B[i] not in graph:
            graph[B[i]] = [A[i]]
        else:
            graph[B[i]].append(A[i])
    print(graph)
    visited = [False] * (N+1)
    visited[0] = True
    visited[1] = True
    queue = [1]
    while queue:
        current = queue.pop(0)
        for i in graph[current]:
            if visited[i] == False:
                visited[i] = True
                queue.append(i)
    return visited[N]

print(solution(4, [1, 2, 4, 4, 3], [2, 3, 1, 3, 1]))
print(solution(4, [1, 2, 1, 3], [2, 4, 3, 4]))
print(solution(6, [2, 4, 5, 3], [3, 5, 6, 4]))
print(solution(3, [1, 3], [2, 2]))

The error it's throwing is :

{1: [2, 4, 3], 2: [1, 3], 3: [2, 4, 1], 4: [1, 3]}
True
{1: [2, 3], 2: [1, 4], 4: [2, 3], 3: [1, 4]}
True
{2: [3], 3: [2, 4], 4: [5, 3], 5: [4, 6], 6: [5]}
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_5588/4101128367.py in <module>
     27 print(solution(4, [1, 2, 4, 4, 3], [2, 3, 1, 3, 1]))
     28 print(solution(4, [1, 2, 1, 3], [2, 4, 3, 4]))
---> 29 print(solution(6, [2, 4, 5, 3], [3, 5, 6, 4]))
     30 print(solution(3, [1, 3], [2, 2]))

~\AppData\Local\Temp/ipykernel_5588/4101128367.py in solution(N, A, B)
     19     while queue:
     20         current = queue.pop(0)
---> 21         for i in graph[current]:
     22             if visited[i] == False:
     23                 visited[i] = True

KeyError: 1

Can anyone help me with this?


Save yourself trouble by using defaultdict . In this problem you only need to check if i-1 is connected to i where i is the number of the node.

from collections import defaultdict
def solution(N, A, B):
    # write your code in Python 3.6
    pass
    graph = defaultdict(set)
    for i in range(len(A)):
        graph[A[i]].add(B[i])
        graph[B[i]].add(A[i])
    print(graph)
    for i in range(2,N+1):
        if i-1 not in graph[i]:
            return False
    return True

print(solution(4, [1, 2, 4, 4, 3], [2, 3, 1, 3, 1]))
print(solution(4, [1, 2, 1, 3], [2, 4, 3, 4]))
print(solution(6, [2, 4, 5, 3], [3, 5, 6, 4]))
print(solution(3, [1, 3], [2, 2]))