Finding missing k numbers in 0…N-1

This is the interview question from http://stackoverflow.com/questions/3492302/easy-interview-question-got-harder-given-numbers-1-100-find-the-missing-numbe?rq=1:

You are given a bag of numbers which are 0 ~ N-1 integers except k numbers. Return the k missing numbers. The solution should operate in O(N) time complexity and O(k) space complexity.

Although the highest voted post uses some theorem to solve the question, I would like to expand another solution which seems more intuitive and practical.

 

Code

class Solution(object):
    def find_k_missing(self, a, n):
        k = n - len(a)
        if k == 0:
            return a
        elif k == n:
            return range(n)

        a.extend([a[0]] * k)

        for i in xrange(n):
            while a[i] != a[a[i]]:
                t = a[i]
                a[i] = a[a[i]]
                a[t] = t

        return [idx for idx, v in enumerate(a) if idx != v]


if __name__ == "__main__":
    s = Solution()
    print s.find_k_missing([3,1,0,4], 5)

 

 

Idea

Suppose we have a bag a which contains the N-k numbers.

1. Extend the a with k additional numbers. The k additional numbers are set as the same value as any existing number, say, a[0]. After this step, a has N numbers.
2. We want to reorder a such that a[j] == j for j in the original N-k numbers.  To do so, traverse every number a[i] in a: swap a[i] and a[a[i]] until a[a[i]] == a[i]. This makes sure at least a[a[i]] == a[i] after the swap, which means a[a[i]] will be associated with the correct number.
3. Then, any position a[j] != j denotes the missing number.

 

 

Reference:

http://pastebin.com/9jZqnTzV