Refactored segment searching method

1 files changed, 25 insertions, 16 deletions

diff --git a/plot.py b/plot.py
index c691d0f..2ae94b9 100644
--- a/plot.py
+++ b/plot.py
@@ -27,29 +27,37 @@ def load(directory):
 	return [(x, _load(y)) for x, y in sorted(_files(directory), key=lambda x: x[0])]
 
 
-def increment(x):
-	return x + 1
+def look_downwards(data, x, start):
+	for i in range(start, 0, -1):
+		if data[i - 1][0] < x:
+			break
+	else:
+		raise IndexError
+	return i - 1
+
 
+def look_upwards(data, x, start):
+	for i in range(start, len(data)):
+		if data[i + 1][0] > x:
+			break
+	else:
+		raise IndexError
+	return i
 
-def decrement(x):
-	return x - 1
+
+def find_segment(data, x):
+	width = data[-1][0] - data[0][0]
+	relative = x - data[0][0]
+	candidate = math.floor(relative / width * len(data))
+	look = look_downwards if data[candidate][0] > x else look_upwards  # May raise IndexError
+	candidate = look(data, x, candidate)
+	return candidate, candidate + 1
 
 
 def find_boundary_curves(swis, x, y):
 	segments = deque()
 	for index, data in swis:
-		width = data[-1][0] - data[0][0]
-		relative = x - data[0][0]
-		candidate = math.floor(relative / width * len(data)) + 1
-		condition = operator.lt if data[candidate][0] > x else operator.gt
-		step = decrement if data[candidate][0] > x else increment
-		j = candidate
-		i = candidate
-		while not condition(data[i][0], x):
-			j = i
-			i = step(i)
-		if i > j:
-			i, j = j, i
+		i, j = find_segment(data, x)
 		if data[i][1] > y and data[j][1] > y:
 			segments.append((index, data, i, j))
 			break
@@ -62,6 +70,7 @@ def find_boundary_curves(swis, x, y):
 
 swis = load("dataset")
 
+
 def onclick(event):
 	if event.button != 1:
 		return