Search a 1d needle in a 1d haystack with numpy

numpy_needle_in_haystack.py

#!/usr/bin/env python

import numpy as np


def np_find_1d(needle, haystack):
    needle = np.asanyarray(needle)
    haystack = np.ascontiguousarray(haystack)

    assert needle.ndim == 1
    assert haystack.ndim == 1

    end = len(haystack) - len(needle)
    view_shape = (end + 1, len(needle))

    search_view = np.ndarray(view_shape, dtype=haystack.dtype, buffer=haystack,
            strides=(haystack.itemsize, haystack.itemsize))

    return np.all(search_view == needle, axis=1)


if __name__ == '__main__':
    a = np.array([5, 3, 1, 9, 1, 5])
    print(a)

    # return True where there is an occurrence
    print(np_find_1d([1, 5], a))

    # return all indices where there is an occurrence
    print(np.where(np_find_1d([1, 5], a)))