Solving it by the logic:

  • find the length
  • traverse to length-n+1
  • remove item

I've written the following code

class ListNode:
 def __init__(self, x):
     self.val = x
     self.next = None

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution:
    # @param A : head node of linked list
    # @param B : integer
    # @return the head node in the linked list
    def removeNthFromEnd(self, A, B):
        temp = A
        len = 0
        while temp is not None:
            len += 1
            temp = temp.next
        result = ListNode(None)
        result.next =  A

        if B >= len:
            if B ==1:
                result.next = None
                result.next = A.next
            for i in range(len-B-1):
                A = A.next
            if B == 1:
                A.next = None
                A.next = A.next.next
        return result.next

if __name__ == "__main__":
        a = ListNode(1)
        b = ListNode(2)
        c = ListNode(3)
        d = ListNode(4)
        e = ListNode(5)
        a.next = b
        b.next = c
        c.next = d
        d.next = e
        A = Solution().removeNthFromEnd(a, 1)
        while A is not None:
            print A.val
            A = A.next

It seems duplicating some logic. How can we improve this code?


Firstly, the logic of the removeNthFromEnd function can be split into these parts:

  • calculate length of list
  • calculate index \$i\$ from the beginning of the list, based on B
  • remove \$i\$th element from the beginning

1. Calculating the length

def list_len(head):
  length = 0
  while head:
    head = head.next
    length += 1
  return length

2. Calculating index from the beginning of the list

i = list_len(A) - B

3. Removing \$i\$th element

def remove_ith(head, i):

  if i == 0:
    return head.next

  lst = head
  for _ in range(i-1):
    lst = lst.next

  lst.next = lst.next.next

  return head

So your removeNthFromEnd function will look like this:

def remove_nth_from_end(head, n):
  length = list_len(head)
  i = length - n

  if not (0 <= i < length):
    return head

  return remove_ith(head, i)

I took the liberty and renamed removeNthFromEnd to obey PEP8, a Python style guide, according to which function names should be in lower_snake_case.

I also renamed parameters from A, B to head, n. Descriptive names with some meaning are always better than just some random letter, even if it's more typing. However sometimes a letter is enough, such as in case of n or i. Especially if the function name is removeNthFromEnd – one would expect such function takes an argument n.

This basic solution has the same performance disadvantage as your solution – in the worst case, you traverse the list twice.

That's not very efficient, at least not time-wise. So I did a little optimization, which in this case might be considered cheating: I convert the linked list into a list of nodes. That way the linked list is only traversed once and you have constant-time access to all nodes. The disadvantage is of course increased memory usage, which is \$O(n)\$ compared to \$O(1)\$ of your solution, where \$n\$ is the length of the list. Optimally you would use a doubly-linked list. That way you would traverse the list at most once (on average only half of the list would be traversed), and the auxiliary space complexity would be constant, as with your original solution.

def remove_nth_from_end_(head, n):

  as_list = []
  node = head

  # Convert linked list to "array".
  while node:
    node = node.next

  # Calculate index from beginning.
  i = len(as_list) - n

  # Check that the list is not empty (in which case nothing can be removed)
  # and the index is within the bounds of the list.
  if as_list and i in range(len(as_list)):

    if i == 0:
      # Removing first element – just shift the head.
      head = head.next

      # Set next-pointer of the previous element to the next element.
      # This is where the constant-time access happens.
      as_list[i-1].next = as_list[i].next

  return head

To further optimize the space complexity, you could only keep the last n+1 elements while traversing the list. That would reduce the auxiliary space complexity to \$O(i_{e})\$, \$i_{e}\$ being the index of the removed element based from the end of the list.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.