Solution for Instagram engineering challenge

unshred.rb

#	USAGE:
#	ruby unshred.rb photo_name_1.png [photo_name_2.png] ...
#
#	it will create an unshredded photo with the original name with "unshredded" prepended
#

require 'RMagick'
include Magick

SHRED_WIDTH = 32
WINDOW_SIZE = 5

#colorspace "distance" between two pixels
def pixel_dist(pix_a, pix_b)
	red = (pix_a.red - pix_b.red)**2
	green = (pix_a.green - pix_b.green)**2
	blue = (pix_a.blue - pix_b.blue)**2

	Math.sqrt(red + green + blue)
end

#average colors for two pixel lists
def pixel_averages(left_pix, right_pix)
	l_red = l_green = l_blue = 0
	r_red = r_green = r_blue = 0
	left_pix.length.times do |i|
		l_red	+= left_pix[i].red
		l_green	+= left_pix[i].green
		l_blue	+= left_pix[i].blue
		r_red	+= right_pix[i].red
		r_green	+= right_pix[i].green
		r_blue	+= right_pix[i].blue
	end

	l_red, l_green, l_blue = [l_red, l_green, l_blue].map! {|c| c /= left_pix.length}
	r_red, r_green, r_blue = [r_red, r_green, r_blue].map! {|c| c /= right_pix.length}

	return Pixel.new(l_red, l_green, l_blue), Pixel.new(r_red, r_green, r_blue)
end

#the total colorspace difference between two strips
def total_diff(left_pix, right_pix, width, height, window_size)
	total = 0
	height.times do |i|
		left_index = width*(i+1) - window_size
		right_index = width*i
		l_ave_pix, r_ave_pix = pixel_averages(left_pix[left_index, window_size], right_pix[right_index, window_size])
		total += pixel_dist(l_ave_pix, r_ave_pix)
	end

	return total
end

#when two strips think they are followed by the same one
#we resolve by reasigning the one with the larger distance
def resolve(keys, diffs, k1)
	s1_min = diffs[k1].min
	keys.values.each do |k2|
		if k1 != k2
			s2_min = diffs[k2].min
			if s1_min[1] == s2_min[1]
				if s1_min[0] < s2_min[0]
					diffs[k2].delete_if {|e| e == s2_min}
					resolve(keys, diffs, k2)
				else
					diffs[k1].delete_if {|e| e == s1_min}
					resolve(keys, diffs, k1)
				end
			end
		end
	end
end


#main
ARGV.each do |a|
	puts "Working on: "+a.inspect
	img = ImageList.new(a)
	img.display
	HEIGHT = img.rows
	shreds = Array.new(img.columns/SHRED_WIDTH)
	keys = {} #so we dont have to use the whole shred for indexing and comparison

	puts "Loading shreds..."
	shreds.length.times do |i|
		shreds[i] = img.get_pixels(i*SHRED_WIDTH, 0, SHRED_WIDTH, HEIGHT)
		keys[shreds[i]] = i
	end

	diffs = {}
	min_diffs = {}

	puts "Getting diffs..."
	keys.values.each do |k1|
		diffs[k1] = []
		keys.values.each do |k2|
			if k1 != k2
				diff = total_diff(keys.key(k1), keys.key(k2), SHRED_WIDTH, HEIGHT, WINDOW_SIZE)
				diffs[k1] << [diff, k2]
			end
		end
	end

	puts "Finding next shreds..."
	#find next shred for each shred
	keys.values.each do |k|
		resolve(keys, diffs, k)
	end

	#get rid of all the data we dont need so lookup on min diff is fast
	keys.values.each do |k|
		diffs[k] = diffs[k].min
	end

	new_shreds = Array.new(shreds.length)
	new_shreds[0] = shreds[0]
	prev_key = 0
	final_keys = []

	#make array of indexes into original shred list for final image
	new_shreds.length.times do |i|
		next_key = diffs[prev_key][1]
		final_keys << prev_key
		prev_key = next_key
		new_shreds[(i+1)%new_shreds.length] = shreds[next_key]
	end

	max_diff = 0
	start_index = 0
	#rotate image so start is in right place
	final_keys.each_with_index do |k, i|
		if max_diff < diffs[k][0]
			max_diff = diffs[k][0]
			start_index = (i+1) % final_keys.length
		end
	end

	start_index.times do
		new_shreds.push(new_shreds.shift)
	end

	new_shreds.length.times do |i|
		img.store_pixels(SHRED_WIDTH*i, 0, SHRED_WIDTH, HEIGHT, new_shreds[i])
	end

	img.display

	img.write("unshredded_" + a)

end