algorithms/remove_nth_from_linked_list.py at master · Merbs/algorithms · GitHub

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

class Node:
    def __init__(self, el, next=None):
        self.el = el
        self.next = next


def noderize(arr):
    head = None
    for el in reversed(arr):
        head = Node(el, head)
    return head

def arrayize(node):
    arr = []
    while True:
        arr.append(node.el)
        if node.next == None:
            break
        node = node.next
    return arr


# remove nth node from linked list
def removeNthFromBeg(head, n):
    if head == None: return None
    if n == 0: return head.next
    initial_head = head
    i = 0
    while head.next:
        if i == n - 1:
            head.next = head.next.next
            break
        i += 1
        head = head.next

    return initial_head


test_cases = [
    ([['A', 'B', 'C', 'D'], 0], ['B', 'C', 'D']),
    ([['A', 'B', 'C', 'D'], 1], ['A', 'C', 'D']),
    ([['A', 'B', 'C', 'D'], 2], ['A', 'B', 'D']),
    ([['A', 'B', 'C', 'D'], 3], ['A', 'B', 'C']),
    ([['A', 'B', 'C', 'D'], 4], ['A', 'B', 'C', 'D'])
]

for test_case, ans in test_cases:
    linked_list, n = test_case
    assert arrayize(removeNthFromBeg(noderize(linked_list), n)) == ans
    assert arrayize(removeNthFromBeg(noderize(linked_list), n)) == ans


# remove nth node from linked list
def removeNthFromEnd(head, n):
    if head == None: return None
    laggard = head
    runner = head
    if n == 0: return head.next
    i = 1
    while runner.next:
        if i > n:
            laggard = laggard.next
        i += 1
        runner = runner.next
    if i == n: return head.next
    if i > n:
        laggard.next = laggard.next.next

    return head


test_cases = [
    ([['A', 'B', 'C', 'D'], 1], ['A', 'B', 'C']),
    ([['A', 'B', 'C', 'D'], 2], ['A', 'B', 'D']),
    ([['A', 'B', 'C', 'D'], 3], ['A', 'C', 'D']),
    ([['A', 'B', 'C', 'D'], 4], ['B', 'C', 'D']),
    ([['A', 'B', 'C', 'D'], 5], ['A', 'B', 'C', 'D'])
]

for test_case, ans in test_cases:
    linked_list, n = test_case
    print linked_list, n, arrayize(removeNthFromEnd(noderize(linked_list), n))
    assert arrayize(removeNthFromEnd(noderize(linked_list), n)) == ans
    assert arrayize(removeNthFromEnd(noderize(linked_list), n)) == ans