Linked List Interview questions

I hadn’t done much programming for a while and thought I’d best brush up on some basic data structures and algorithm.

The following is from Programming Interviews Exposed. I implemented some of the Linked List stuff in the first section in Python. The code came out pretty close to the solutions in the book which were mainly in C++.

class Stack:
    """
    Adapted from Programming Interviews Exposed problems
    LinkedList implementation that keeps track of head and tail

    Basic tests - push,pop,remove
    >>> s=Stack()
    >>> s.push(1)
    >>> s.push(2)
    >>> s.push(3)
    >>> s.push(4)
    >>> s.pop()
    4
    >>> e=s.head.next
    >>> s.remove(e)
    True
    >>> s.head.data
    3
    >>> s.tail.data
    1
    >>> e=s.head.next
    >>> s.remove(e)
    True
    >>> s.head.data
    3
    >>> s.tail.data
    3
    >>> e=s.head
    >>> s.remove(e)
    True
    >>> s.head
    >>> s.tail
    """

    def __init__(self):
        self.head=None
        self.tail=None

    def push(self,data):
        if self.head==None:
            self.head=Element(data,None)
            self.tail=self.head
        else:
            self.head=Element(data,self.head)

    def pop(self):
        oldhead=self.head
        if oldhead==None:
            return None
        self.head=oldhead.next
        if self.head==None:
            self.tail=None
        return oldhead.data

    def remove(self,element):
        if self.head==None:
            return False
        if self.head==element:
            self.head=self.head.next
            if self.head==None:
                self.tail=None
            return True
        curr=self.head
        while curr!=None:
            if curr.next==element:
                curr.next=curr.next.next
                if curr.next==None:
                    self.tail=curr
                return True
            else:
                curr=curr.next
                return False

    def insertAfter(self,element,data):
        """
        >>> s=Stack()
        >>> s.insertAfter(None,1)
        True
        >>> s.head.data
        1
        >>> s.tail.data
        1
        >>> s.insertAfter(None,2)
        True
        >>> s.head.data
        2
        >>> s.tail.data
        1
        >>> e=s.head
        >>> s.insertAfter(e,3)
        True
        >>> s.head.data
        2
        >>> s.tail.data
        1
        >>> s.head.next.data
        3
        >>> e=s.tail
        >>> s.insertAfter(e,4)
        True
        >>> s.head.data
        2
        >>> s.tail.data
        4
        """
        if (element==None)|(self.head==None):
            self.head=Element(data,self.head)
            if self.head.next==None:
                self.tail=self.head
            return True
        curr=self.head
        while curr!=None:
            if curr==element:
                curr.next=Element(data,curr.next)
                if curr.next.next==None:
                    self.tail=curr.next
                return True
            curr=curr.next
        return False

    def removeHead(self):
        """
        >>> s=Stack()
        >>> s.removeHead()
        >>> s.head
        >>> s.tail
        >>> s.push(1)
        >>> s.removeHead()
        >>> s.head
        >>> s.tail
        >>> s.push(1)
        >>> s.push(2)
        >>> s.removeHead()
        >>> s.head.data
        1
        >>> s.tail.data
        1
        """
        if self.head!=None:
            self.head=self.head.next
        if self.head==None:
            self.tail=None

    def findMToLastElement(self,m):
        """
        >>> s=Stack()
        >>> s.push(1)
        >>> s.push(2)
        >>> s.push(3)
        >>> s.findMToLastElement(0).data
        1
        >>> s.findMToLastElement(1).data
        2
        >>> s.findMToLastElement(2).data
        3
        >>> s.findMToLastElement(3)
        """
        mToLast=None
        curr=self.head
        count=0
        while curr!=None:
            if count==m:
                mToLast=self.head
            if count>m:
                mToLast=mToLast.next
                if mToLast==None:
                    return None
            curr=curr.next
            count+=1
        return mToLast

class Element:
    def __init__(self,data,next=None):
        self.data=data
        self.next=next

if __name__ == "__main__":
    import doctest
    doctest.testmod()