Use this to learn the idea, then write your own version.
123 4import random5import itertools6import math7 8 91011random.seed(0)12class Solution(object):13 def beautifulPair(self, nums1, nums2):14 """15 :type nums1: List[int]16 :type nums2: List[int]17 :rtype: List[int]18 """19 INF = float("inf")20 def dist(a, b):21 if a[2] > b[2]:22 a, b = b, a23 return [abs(a[0]-b[0])+abs(a[1]-b[1]), a[2], b[2]]24 25 def cell(point, size):26 x, y, _ = point27 return math.floor(x/size), math.floor(y/size)28 29 def improve():30 lookup = {}31 for p in points:32 i, j = map(int, cell(p, result[0]/2.0))33 for ni in xrange(i-2, (i+2)+1):34 for nj in xrange(j-2, (j+2)+1):35 if (ni, nj) not in lookup:36 continue37 d = dist(p, lookup[ni, nj])38 if d < result:39 result[:] = d40 return True41 lookup[i, j] = p42 return False43 44 points = [(i, j, idx) for idx, (i, j) in enumerate(itertools.izip(nums1, nums2))]45 result = [INF]*346 lookup = {}47 for i in reversed(xrange(len(points))):48 if points[i][:2] in lookup:49 result = [0, i, lookup[points[i][:2]]]50 lookup[points[i][:2]] = i51 if result[0] == 0:52 return result[1:]53 random.shuffle(points)54 result = dist(points[0], points[1])55 while improve():56 pass57 return result[1:]58 59 606162import itertools63 64 656667class Solution2(object):68 def beautifulPair(self, nums1, nums2):69 """70 :type nums1: List[int]71 :type nums2: List[int]72 :rtype: List[int]73 """74 INF = float("inf")75 MAX_NEIGHBOR_COUNT = (8+2)276 def dist(a, b):77 if a > b:78 a, b = b, a79 return [abs(points[a][0]-points[b][0])+abs(points[a][1]-points[b][1]), a, b]80 81 def merge_sort(left, right):82 def update(arr, i): 83 for j in reversed(xrange(len(arr))):84 if points[i][1]-points[arr[j]][1] > result[0]:85 break86 result[:] = min(result, dist(i, arr[j]))87 else:88 j = -189 assert((len(arr)-1)-j <= MAX_NEIGHBOR_COUNT)90 91 if left == right:92 return93 mid = left+(right-left)294 x = points[order[mid]][0] 95 merge_sort(left, mid)96 merge_sort(mid+1, right)97 tmp, tmp_l, tmp_r = [], [], []98 l, r = left, mid+199 while l <= mid or r <= right:100 if r == right+1 or (l <= mid and points[order[l]][1] <= points[order[r]][1]): 101 update(tmp_r, order[l])102 if x-points[order[l]][0] <= result[0]: 103 tmp_l.append(order[l])104 tmp.append(order[l])105 l += 1106 else:107 update(tmp_l, order[r])108 if points[order[r]][0]-x <= result[0]: 109 tmp_r.append(order[r])110 tmp.append(order[r])111 r += 1112 order[left:right+1] = tmp113 114 points = [(i, j) for i, j in itertools.izip(nums1, nums2)]115 result = [INF]*3116 lookup = {}117 for i in reversed(xrange(len(points))):118 if points[i] in lookup:119 result = [0, (i, lookup[points[i]])]120 lookup[points[i]] = i121 if result[0] == 0:122 return result[1]123 order = range(len(points))124 order.sort(key=lambda x: points[x][0])125 merge_sort(0, len(points)-1)126 return result[1:]127 128 129130131import itertools132 133 134135136class Solution3(object):137 def beautifulPair(self, nums1, nums2):138 """139 :type nums1: List[int]140 :type nums2: List[int]141 :rtype: List[int]142 """143 INF = float("inf")144 MAX_NEIGHBOR_COUNT = 8145 def dist(a, b):146 if a > b:147 a, b = b, a148 return [abs(points[a][0]-points[b][0])+abs(points[a][1]-points[b][1]), a, b]149 150 def merge_sort(left, right):151 if left == right:152 return153 mid = left + (right-left)2154 x = points[order[mid]][0] 155 merge_sort(left, mid)156 merge_sort(mid+1, right)157 r = mid+1158 tmp = []159 for l in xrange(left, mid+1):160 while r <= right and points[order[r]][1] < points[order[l]][1]: 161 tmp.append(order[r])162 r += 1163 tmp.append(order[l])164 order[left:left+len(tmp)] = tmp165 166 167 stripe = [order[i] for i in xrange(left, right+1) if abs(points[order[i]][0]-x) <= result[0]]168 for i in xrange(len(stripe)-1):169 for j in xrange(i+1, len(stripe)):170 x, y = stripe[i], stripe[j]171 if points[y][1]-points[x][1] > result[0]:172 break173 result[:] = min(result, dist(x, y))174 else:175 j = len(stripe)176 assert(j-(i+1) <= MAX_NEIGHBOR_COUNT)177 178 points = [(i, j) for i, j in itertools.izip(nums1, nums2)]179 result = [INF]*3180 lookup = {}181 for i in reversed(xrange(len(points))):182 if points[i] in lookup:183 result = [0, (i, lookup[points[i]])]184 lookup[points[i]] = i185 if result[0] == 0:186 return result[1]187 order = range(len(points))188 order.sort(key=lambda x: points[x][0])189 merge_sort(0, len(points)-1)190 return result[1:]191 192 193194195import itertools196 197 198199class Solution4(object):200 def beautifulPair(self, nums1, nums2):201 """202 :type nums1: List[int]203 :type nums2: List[int]204 :rtype: List[int]205 """206 INF = float("inf")207 208 class SegmentTree(object):209 def __init__(self, N,210 build_fn=lambda _: [-INF, -INF], 211 query_fn=lambda x, y: y if x is None else x if y is None else max(x, y),212 update_fn=lambda x: x):213 self.tree = [None]*(2*2**((N-1).bit_length()))214 self.base = len(self.tree)2215 self.query_fn = query_fn216 self.update_fn = update_fn217 for i in xrange(self.base, self.base+N):218 self.tree[i] = build_fn(i-self.base)219 for i in reversed(xrange(1, self.base)):220 self.tree[i] = query_fn(self.tree[2*i], self.tree[2*i+1])221 222 def update(self, i, h):223 x = self.base+i224 self.tree[x] = self.update_fn(h)225 while x > 1:226 x = 2227 self.tree[x] = self.query_fn(self.tree[x*2], self.tree[x*2+1])228 229 def query(self, L, R):230 if L > R:231 return [-INF, -INF] 232 L += self.base233 R += self.base234 left = right = None235 while L <= R:236 if L & 1:237 left = self.query_fn(left, self.tree[L])238 L += 1239 if R & 1 == 0:240 right = self.query_fn(self.tree[R], right)241 R -= 1242 L = 2243 R = 2244 return self.query_fn(left, right)245 246 def dist(a, b):247 if a > b:248 a, b = b, a249 return [abs(points[a][0]-points[b][0])+abs(points[a][1]-points[b][1]), a, b]250 251 points = [(i, j) for i, j in itertools.izip(nums1, nums2)]252 result = [INF]*3253 lookup = {}254 for i in reversed(xrange(len(points))):255 if points[i] in lookup:256 result = [0, (i, lookup[points[i]])]257 lookup[points[i]] = i258 if result[0] == 0:259 return result[1]260 order = range(len(points))261 order.sort(key=lambda x: points[x][0])262 y_set = set(y for _, y in points)263 y_to_idx = {y:i for i, y in enumerate(sorted(y_set))}264 st1, st2 = SegmentTree(len(y_to_idx)), SegmentTree(len(y_to_idx))265 for i in order:266 j = -st1.query(0, y_to_idx[points[i][1]]-1)[1] 267 if j != INF:268 assert(points[j][1] < points[i][1])269 result = min(result, dist(i, j))270 st1.update(y_to_idx[points[i][1]], [points[i][0]+points[i][1], -i])271 j = -st2.query(y_to_idx[points[i][1]], len(y_to_idx)-1)[1] 272 if j != INF:273 assert(points[j][1] >= points[i][1])274 result = min(result, dist(i, j))275 st2.update(y_to_idx[points[i][1]], [points[i][0]-points[i][1], -i])276 return result[1:]277