grid-segmentation

grid-segmentation.cpp

#include <iostream>
#include <vector>
#include <map>
#include <set>

using namespace std;

const int M = 9;
const int N = M * M;
const int S = 9;
const int MAX_COUNT = N / S;

set<int> neighbors(const vector<int>& grid, int idx)
{
    set<int> b;
    int x = idx % M;
    int y = idx / M;
    if (x - 1 >= 0)
        b.insert(grid[idx - 1]);
    if (x + 1 < M)
        b.insert(grid[idx + 1]);
    if (y - 1 >= 0)
        b.insert(grid[idx - M]);
    if (y + 1 < M)
        b.insert(grid[idx + M]);
    b.erase(0);
    return b;
}

void print_grid(vector<int> grid)
{
    for (int i = 0; i < N; i++) {
        int x = i % M;
        cout << grid[i];
        if (x == M - 1)
            cout << endl;
    }
    cout << endl;
}

int total = 0;

void put(vector<int> grid, int idx, vector<int> count)
{
    if (idx == N) {
        total++;
        print_grid(grid);
        return;
    }
    auto b = neighbors(grid, idx);
    set<int> alter;
    // check if full
    for (int i = 1; i <= S; i++) {
        if (count[i] == MAX_COUNT)
            continue;
        alter.insert(i);
    }
    
    for (auto a : alter) {
        // check if continuous 
        if (count[a] > 0 && b.find(a) == b.end())
            continue;
        auto next = grid;
        auto nc = count;
        next[idx] = a;
        nc[a]++;
        put(next, idx + 1, nc);
    }
}

int main()
{
    vector<int> grid(N);
    vector<int> count(S+1);
    put(grid, 0, count);
    cout << "total:" << total << endl;
}