Write nested loops as linq query

Question

I am tracking trains and trying to identify individual trains seen multiple times at different points through the IDs of the wagons on them when spotted.

// create a lookup of all tracked trains for each wagon
IEnumerable<Train> trains = GetTrains();
var wagonTrains = new Dictionary<int, List<Train>>();
foreach (Train t in trains)
{
 foreach (int w in t.WagonsInTrain)
 {
 if (!wagonTrains.ContainsKey(w))
 {
 wagonTrains.Add(w, new List<Train>());
 }
 wagonTrains[w].Add(t);
 }
}

Is there a better way do to what I am doing in this code segment? Perhaps some chained linq operations? Is there a particular name for the kind of operation I am using here?

Answer 1

create a lookup of all tracked trains for each wagon

That's pretty much what the ToLookup() method is for. You just need a bit more LINQ to get a collection of (train, wagon) pairs, so that ToLookup() can work:

var wagonTrains =
 (from train in trains
 from wagon in train.WagonsInTrain
 select new { train, wagon })
 .ToLookup(x => x.wagon, x => x.train);

Answer 2

Here is one that uses chained linq operations:

var trainsByWagon =
 trains
 .SelectMany(train => train.WagonsInTrain, (train, wagon) => new { train, wagon })
 .GroupBy(trainAndWagon => trainAndWagon.wagon, trainAndWagon => trainAndWagon.train);
 .ToDictionary(g => g.Key, g => g.ToList());

Small print: Haven't actually tried this, but should work.

• Michael Sandler: Sorry, I eventually accepted the chronologicaly first answer, but both are super useful.

Answer 3

I've actually set this up in my code as a pair of extension methods:

public static class IDictionaryExtensions
{
 public static void Update<TKEY, TVALUE>(this IDictionary<TKEY, TVALUE> dictionary, TKEY key, TVALUE value)
 {
 if (dictionary.ContainsKey(key))
 {
 dictionary[key] = value;
 }
 else
 {
 dictionary.Add(key, value);
 }
 }
 public static void AddToList<TKEY, TVALUE>(this IDictionary<TKEY, IList<TVALUE>> dictionary, TKEY key, TVALUE value)
 {
 if (dictionary.ContainsKey(key) && dictionary[key] != null)
 {
 dictionary[key].Add(value);
 }
 else
 {
 dictionary.Update(key, new List<TVALUE> {value});
 }
 }
 }

The first method is an all-purpose "Add"-type method which will either add it or update it depending on whether or not it exists. The second is specifically for dictionaries where the value is a List<>. It'll let you add an element to the List<>, creating the key if neccesary.

---

For your example, usage would be:

IEnumerable<Train> trains = GetTrains();
var wagonTrains = new Dictionary<int, List<Train>>();
foreach (var train in trains)
{
 foreach (var wagon in train.WagonsInTrain)
 {
 wagonTrains.AddToList(wagon, train);
 }
}

• svick: Your code wouldn't compile, AddToList() has two parameters.
• Bobson: @svick - Good catch. That's what I get for freehanding it. Fixed.