Van_der_Waerden.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;

namespace ConsoleApplication6 {
    class Program {
        static void Main(string[] args) {
            // cold start for JIT
            Time(2, 3);

            // hot times
            Time(2, 3); // 9
            Time(2, 4); // 35
            Time(2, 5); // 178
            //Time(2, 6); // 1132
            Time(3, 3); // 27
            //Time(3, 4); // 293
            Time(4, 3); // 76
        }

        static void Time(byte r, int k) {
            GC.Collect(3, GCCollectionMode.Forced);
            var sw = Stopwatch.StartNew();
            var result = new Van_der_WaerdenFinder(r, k).Invoke();
            sw.Stop();
            Console.WriteLine("W({0}, {1}) = {2} (elapsed: {3})", r, k, result, sw.Elapsed);
        }
    }

    // ReSharper disable once InconsistentNaming
    public sealed class Van_der_WaerdenFinder {
        private readonly byte _maxcolor;

        private readonly int _progressionLength;

        // a depth first search avoids holding too many sequences in memory at the same time
        // because the odds that a particular color choice is one we want to keep looking for is
        // roughly 1/x where x is the number of colors in sequence after the one we are currently selecting
        private readonly ConcurrentStack<Frame> _collection;

        private readonly Action _worker;
        private int _lowerBound;
        private readonly byte _r;

        struct Frame {
            public readonly ImmutableStack<byte> Stack;
            public readonly int Count;

            public Frame(ImmutableStack<byte> stack, int count) {
                Stack = stack;
                Count = count;
            }
        }

        public Van_der_WaerdenFinder(byte r, int k) {
            _r = r;
            _maxcolor = (byte) ((1 << r) - 1);
            _progressionLength = k;
            _collection = new ConcurrentStack<Frame>();
            _lowerBound = 0;
            _worker = Worker;
        }

        public int Invoke() {
            // basic idea: spawn sufficient worker threads to consume available CPU resources
            // on main thread produce a few trivial sequences (for r=2,k=3: [1,2,1], [1,2,2], [1,1,2,1], [1,1,2,2]; add for k=4: [1,1,1,2,1], [1,1,1,2,2] and so on)
            // to increase probability that no workers immediatly quit
            _collection.PushRange(Start().ToArray());

            // then create worker threads and wait for them to finish
            // todo: improve this to ensure workers don't immediately quit
            var workers = new List<Task>();
            for (int i = 0; i < Environment.ProcessorCount * 2; i++) {
                workers.Add(Task.Run(_worker));
            }
            Task.WaitAll(workers.ToArray());

            // at this point an exhaustive search has been completed for all sequence lengths `_lowerBound + 1` and no such length sequences have been produced
            // therefore this is the number we are looking for
            return _lowerBound + 1;
        }

        void Worker() {
            Frame item;
            // there is a slight amount of thread synchronization in producing/consuming the top item from the shared stack of sequences
            // to avoid that, this double loop runs depth first along the first item in the production (generally following along the color "1")
            // and doesn't pop off the shared pool until it has followed this depth first traversal as far as it can
            while (_collection.TryPop(out item)) {
                while (true) {
                    var list = Next(item);
                    if (list.IsEmpty) {
                        while (Math.Max(_lowerBound, item.Count) > Interlocked.CompareExchange(ref _lowerBound, Math.Max(_lowerBound, item.Count), _lowerBound)) {
                            // spin ... (item.Count is a better lower bound than the existing value, and it couldn't set due to contention)
                            // This actually sets the value twice when it works correctly, but it happens rarely enough that it doesn't matter for perf
                            // for a quick progress indicator uncomment:
                            //Console.WriteLine("W({0}, {1}) > {2}", _r, _progressionLength, Math.Max(_lowerBound, item.Count));
                            // (for everything uncommented in Main except for 4,3 this gives the correct lb instantly so it is kinda pointless)
                        }
                        break;
                    }
                    list = list.Pop(out item);
                    while (!list.IsEmpty) {
                        Frame value;
                        list = list.Pop(out value);
                        _collection.Push(value);
                    }
                }
            }
        }

        ImmutableStack<Frame> Start() {
            var result = ImmutableStack<Frame>.Empty;
            var ones = ImmutableStack<byte>.Empty;
            var length = _progressionLength - 1;
            for (int i = 0; i < length; i++) {
                ones = ones.Push(1);
                var next = Next(new Frame(ones.Push(2), i + 2));
                foreach (var item in next) {
                    result = result.Push(item);
                }
            }
            return result;
        }

        ImmutableStack<Frame> Next(Frame existing) {
            var result = ImmutableStack<Frame>.Empty;
            var nextColors = GetNextColors(existing);
            for (byte i = 1; i < _maxcolor; i *= 2) {
                if ((nextColors & i) == i) result = result.Push(new Frame(existing.Stack.Push(i), existing.Count + 1));
            }
            return result;
        }


        byte GetNextColors(Frame sequence) {
            // sequence has no progressions of k length where all values are the same number
            // determine byte representing colors that can be added to the list without causing a progression
            var colors = _maxcolor;
            var count = sequence.Count;
            // trivial case
            if (count + 1 < _progressionLength) {
                if (_r > 3) {
                    //optimization for 4+ colors: only unlock 1 new color at a time to avoid palette swapped trees
                    byte usedcolors = 0;
                    var s = sequence.Stack;
                    do {
                        byte color;
                        s = s.Pop(out color);
                        usedcolors |= color;
                    } while (!s.IsEmpty);
                    return (byte) ((usedcolors * 2 + 1) & colors);
                }
                return colors;
            }

            // copying sequence to array for indexing (could instead use an ImmutableList<byte> for the stack...
            // but that is O(log(n)) to index in the loops below and append to in Next() and Start() which costs at least as much as this)
            var array = new byte[count];
            var stack = sequence.Stack;
            var index = 0;
            do {
                byte color;
                stack = stack.Pop(out color);
                array[index] = color;
                index++;
            } while (!stack.IsEmpty);

            if (_r > 3) {
                //optimization for 4+ colors again
                byte usedcolors = 0;
                foreach (byte color in array) {
                    if (usedcolors == colors) break;
                    usedcolors |= color;
                }
                colors = (byte) ((usedcolors * 2 + 1) & colors);
            }

            for (var width = 1; width <= array.Length; width++) {
                if (colors == 0) break;
                index = width - 1;
                var checkvalue = array[index];
                if ((colors & checkvalue) == 0) continue;
                for (int multiplier = _progressionLength - 1; multiplier >= 1; multiplier--) {
                    if (multiplier == 1) {
                        colors -= checkvalue;
                        break;
                    }
                    index = width * multiplier - 1;
                    if (index >= array.Length) goto Complete; //common CLR optimization: don't return from inside loops
                    if (array[index] != checkvalue) break;
                }
            }
            Complete:
            return colors;
        }
    }
}