Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save aodhan-domhnaill/f00e4fbcde3cad68cedb6ff116af6f5d to your computer and use it in GitHub Desktop.
Save aodhan-domhnaill/f00e4fbcde3cad68cedb6ff116af6f5d to your computer and use it in GitHub Desktop.
import numpy as np
from scipy.sparse import dia_matrix
import timeit
def matches(a, b):
mask = dia_matrix((np.ones((1, a.size)).repeat(4, axis=0), np.arange(4)),
shape=(a.size, b.size), dtype=np.int)
matches = a[mask.T.dot(mask.dot(a == b) == 4).astype(np.bool)]
if __name__ == '__main__':
for p in [3, 5, 7]:
print('1e%d -' % p, timeit.timeit('matches(np.arange(1e%d), np.arange(1e%d))' % (p, p),
setup="from __main__ import matches, np, dia_matrix", number=100))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment